ktcapraz/t5gemma2-r1s4-steprl

T5Gemma2 r1s4: decisions that reason with an evidence graph

A research model. Before it answers, it builds a small evidence graph: it copies spans from the text, links them, and decides each link's relation from a fixed list. Code then executes the graph, so the final answer goes through the reasoning by construction. The last training stage is step-level reinforcement learning (GRPO).

It is built from a clean google/t5gemma-2-270m-270m base. It answers choice, noul and score questions.

The idea works on some tasks, but this is not the best way to get accurate answers. On generator rows with arithmetic, reasoning adds 6.7 points. On real-use questions, the final answer (0.633) is worse than the model's own base answer (0.714). If you want the strongest decisions from this family, use t5gemma2-decision-270m-v2, or set reasoning to off here.

The request format follows TypeSafe's System One API. This model is not made by, or affiliated with, TypeSafe.

A reasoning trace

From a real run. The base answer for "Does this require a refund?" is 0.764:

State: Ordered headphones on May 2. They arrived 12 days ago and the left ear stopped working.
       Policy: refunds for defective items within 30 days of delivery. Sale items are final. I used a 20% coupon.

link 1  "the left ear stopped working"  MEETS  "defective items"                           p 0.9999
link 2  "arrived 12 days ago and the left ear stopped working. Policy: refunds for
         defective items within 30 days of delivery"  SUPPORTS  [Yes]                      p 1.0

final logit = base logit + push (Yes +4.59) + correction (Yes -0.05, No +0.06)
Answer: yes 0.997

Spans are positions in the input, so a quote cannot contain invented words. Link 2 shows a known weakness: the span runs across a sentence edge.

Results

Scorecard at the end of training (results/final_scorecard.txt). base is the answer before reasoning. final is the graph answer. off is the base plus the correction, with no links.

Set	final	base	off
generator (code-made rules and arithmetic)	0.837	0.770	0.765
gold_b (rule chains and multi-hop)	0.881	0.867	0.798
contrast pairs	0.773	0.776	0.772
gold_a (gold evidence and links)	0.728	0.733	0.668
plain decisions	0.704	0.710	0.701
tasksource (many tasks)	0.679	0.680	0.679
views (same question, other forms)	0.657	0.669	0.656
real-use (147 q)	0.633	0.714	0.612

On real-use, reasoning rescued 3 wrong answers and damaged 15 right ones, with links on 30% of questions.

S1MB: r1s4 was not scored on the full benchmark. Its parent r1 scored 39.3 (final) and 43.7 (base; trained tasks 51.0, never-trained 25.6). On a 2,825-question sample, the r1s4 base scored 44.9 overall and 31.8 on never-trained tasks; v2 scored 54.5 and 30.4 on the same rows. This sample is not an official score. At the served stop bias, r1s4 linked only 10 of those 2,825 questions.

Quick start

The model needs the code in decisions/, and the base model for the frozen word table and tokenizer. Accept the Gemma terms on the base model page first.

pip install -r requirements.txt
python run_example.py                  # reasoning on
python run_example.py --reasoning off  # the base answer only (fast mode)

In your own code:

from decisions.format import parse_request
from decisions.r1_engine import R1Engine

engine = R1Engine(".", "google/t5gemma-2-270m-270m", threads=6)
body = {
    "state": "Ordered headphones on May 2. They arrived 12 days ago and the left ear stopped working. "
             "Policy: refunds for defective items within 30 days of delivery. Sale items are final. I used a 20% coupon.",
    "questions": {"refund": {"type": "noul", "instructions": "Does this require a refund?"}},
}
state, _ = parse_request(body)
answers, usage, meta = engine.decide(state, body["questions"], reasoning=True)
answer = answers["refund"]
print(answer["noul"], answer["base"]["noul"], answer["reasoning"]["steps"])

Output of the code above: final 0.996524, base 0.764275, with the 2 steps shown in the trace.

Speed on a server CPU with 4 threads, on short texts: about 0.2 to 0.4 s per request with reasoning on, and about 0.2 s with it off.

How it works

Encoder (non-joint). Layers 1 to 14 read the state once, and the state never sees a question. Before layer 15, the state is copied once per question, and each copy reads only its own question. So questions are independent by construction, and one encoding is shared.

Decoder, one row per question:

<bos> [BASE] <think> [link 1] [link 2] ... </think> Answer:

Links. A link is [A] [B] [REL]. The ends are spans (picked by a pointer), options, or earlier links. The relation is one of 9: SUPPORTS, CONTRADICTS, IMPLIES, MEETS, FAILS, EXCEPTION, SAME, COMPARE, NONE. A CALC step does arithmetic in code, with a unit table. If the same pair is linked twice, the last link wins, so NONE can take a link back.

Graph program. Differentiable code computes a strength for each node from the links. Final logit = base logit + logit of the option's strength + a bounded correction r (+-2). With no links, final = base + r. COMPARE and CALC use code when both sides parse as numbers, dates or quantities.

Stop bias. The model uses a stop bias of +3.0 on the decision to stop linking, chosen on dev sets. Lower values give more links and worse answers.

What we learned by examining the model

• The base answer comes from the encoder. With a constant decoder state, accuracy does not change (new tasks 0.438, multi-link 0.867).
• Links mostly repeat the base. The strongest link agrees with the base top answer in 96% of multi-link and 78% of generator cases.
• Real rescues need code tools. On generator rows, 10 of 10 disagreements with the base were right, all with CALC or COMPARE. Without code, 2 of 4 were right on multi-link rows.
• Relations barely read the sentence. With a random wrong sentence as A, P(gold relation) falls only from 0.999 to 0.966.
• It does not reason on new tasks. It wrote no link on 160 new-task questions. Forcing links there gives empty ones.
• The correction hurts on real-use. It pulls down the top base option in 68-82% of questions. Base 0.714, base + correction 0.633, base + links 0.694.
• A cheap fix exists, and it is not applied here. Zeroing ro_r.proj_a and ro_r.bias_o (328,833 weights, 0.088%) lifted the final real-use score from 0.626 to 0.694 in our test (147 q). The released weights are unchanged.

Limitations

• Probabilities are not calibrated. They are plain softmax values, and a link can push an answer by up to +4.59 logits. Do not read them as true likelihoods.
• The final answer is often worse than the base answer on plain tasks and real use.
• Weak on tasks it never trained on.
• Sensitive to word form, as the rest of this model family is.
• Tested on English only.
• No safety evaluation. Do not use it alone for decisions about people, or in medical, legal or financial settings.

Training

Stage	What	Budget
r1 phase 1	Supervised: gold, code-made and teacher links, from the clean base
r1 phase 1.5	Self-training: the model's own graphs that pass the program and counterfactual checks; repair rows from its own mistakes
r1 total	One A100-80GB job	about $25
r1s4	Step-level GRPO from the step 16,000 EMA, through r1s2 step 16,100; learning rate scale 0.25; 20 min; weights are the EMA at step 16,316	$1.62

Questions in the mix: plain decisions 581K (hotchpotch/bekko-system-one-dataset-v0 train and others), gold evidence and links 373K (for example FEVER, VitaminC, ProofWriter, EntailmentBank, ShARC, HotpotQA, MuSiQue), tasksource 198K, intervention 118K, teacher 103K, generator 78K, contrast 49K.

Teacher links and contrast edits came from Claude Haiku 5.5, and were kept only when the graph program computed the gold answer. Training data licences differ between sources, and some are non-commercial. Check these before commercial use.

Files

File	Content
r1s4.safetensors	Trained weights only (204.7M parameters). The frozen word table comes from the base model.
r1_config.json	Model settings, stop bias, step
decisions/	Inference code, including the graph program in decisions/r1/graph_program.py
run_example.py, requirements.txt	Example and pinned versions
results/final_scorecard.txt	End-of-training scorecard
NOTICE	Gemma notice and list of changes

License

Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms. This model is a modified version of google/t5gemma-2-270m-270m: the weights were changed by fine-tuning, and new heads and tokens were added. The Gemma Prohibited Use Policy applies to this model and to anything built from it. The code in decisions/ is provided under the same terms as the model.