oa-rlaif-e4b-grpo-weights

Stage 2, GRPO with the DRAL-OA reward and the A4B judge (RLAIF): a LoRA adapter for base + OPSD, merged. Weights (a LoRA adapter): it needs the base model.

OA-RLAIF aligns ontologies: for each class of a source ontology it picks the corresponding class of a target ontology among retrieved candidates, or NONE. Base model: Gemma 4 E4B. Code: github.com/AGR19/OA-RLAIF (private).

The OA-RLAIF repos

Repo	Kind	What
oa-rlaif-e4b-sft-weights	weights (LoRA)	Stage 1, SFT: adapter for google/gemma-4-E4B-it
oa-rlaif-e4b-opsd-weights	weights (LoRA)	Stage 1, OPSD: adapter for google/gemma-4-E4B-it
oa-rlaif-e4b-grpo-weights	weights (LoRA)	Stage 2, GRPO: adapter for base + OPSD (merged)
oa-rlaif-e4b-sft-model-safetensors	model (safetensors, bf16)	base + SFT
oa-rlaif-e4b-opsd-grpo-model-gguf	model (GGUF Q4_K_M)	base + OPSD + GRPO
oa-rlaif-e4b-all-combined-model-safetensors	model (safetensors, bf16)	base + OPSD → GRPO + SFT
oa-rlaif-e4b-all-combined-model-gguf	model (GGUF Q4_K_M)	base + OPSD → GRPO + SFT, quantized
oa-rlaif-datasource-qa-data	data	OM datasource passages → questions and answers
oa-rlaif-opsd-sft-data	data	Stage 1 rows: the same rows train OPSD and SFT
oa-rlaif-grpo-data	data	Stage 2 prompts with their episodes; evaluation sets; splits
oa-rlaif-dral-reward-data	data	DRAL-OA vectors for the GRPO reward term r_DRAL
oa-rlaif-icd-bridge-data	data	SNOMED CT → ICD-10 → ICD-11 composed links
oa-rlaif-raw-data	data	the raw ontology folders (no PDFs), Bio-ML, parsed-ontology cache
oa-rlaif-heldout-eval-results	results	the earlier evaluation (held-out decisions)
oa-rlaif-full-eval-results	results	the full-alignment evaluation (full_eval/)

How to use

This adapter was trained on base + OPSD, not on the plain base. Merge OPSD first:

import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor

base = AutoModelForImageTextToText.from_pretrained("google/gemma-4-E4B-it", dtype=torch.bfloat16)
opsd = PeftModel.from_pretrained(base, "gonibeed/oa-rlaif-e4b-opsd-weights").merge_and_unload()
model = PeftModel.from_pretrained(opsd, "gonibeed/oa-rlaif-e4b-grpo-weights")
tokenizer = AutoProcessor.from_pretrained("google/gemma-4-E4B-it").tokenizer

(adapter_config.json names the local folder of the merged OPSD model it was trained on.) The merged and quantized result is oa-rlaif-e4b-opsd-grpo-model-gguf.

How the OA-RLAIF models are trained

Stage	Starts from	Data	Loss
SFT (Stage 1)	google/gemma-4-E4B-it	oa-rlaif-opsd-sft-data	cross-entropy on the reference answer y* (off-policy)
OPSD (Stage 1)	google/gemma-4-E4B-it	oa-rlaif-opsd-sft-data (the same rows)	the student samples a reply y; a frozen copy of the base model, also shown the privileged context c (the gold answer and its evidence), scores every token; loss = Σ_t KL(student ‖ teacher)
GRPO (Stage 2)	base + OPSD, merged	oa-rlaif-grpo-data, oa-rlaif-dral-reward-data	group-relative policy gradient with the reward below

All three train LoRA adapters only. LoRA r = 32, α = 64, dropout 0, on the language model's attention and MLP projections (q, k, v, o, gate, up, down); the vision and audio towers stay frozen and unused. bf16. Trained with thinking off (the chat template's enable_thinking=False). All combined = base + OPSD → GRPO (merged), then the SFT adapter's weight change added on top (task arithmetic: the SFT adapter was trained on the plain base).

Training settings

Setting	SFT	OPSD	GRPO
Learning rate, schedule	1e-4, cosine, 3% warm-up (as OPSD)	1e-4, cosine, 3% warm-up	5e-6, linear decay (trainer default), no warm-up
Batch per GPU × accumulation	8 × 4	8 × 4	8 × 2, G = 8 replies per prompt
Completion length	384 tokens	384 tokens (one sample per prompt, temperature 1.0)	256 tokens (temperature 1.0)
Other	completion-only loss	teacher = frozen base with the privileged context; top-100 logit distillation	Dr. GRPO, β = 0, rewards not scaled
Steps (from each adapter's training_args.bin)	437	539	1,068

Hardware: 4 × B200 on Vast.ai (3 GPUs train; 1 GPU serves vLLM rollouts and, in GRPO, the A4B judge).

Prompt and reply format

The model decides one source class at a time from k numbered candidates (40 for Anatomy, 20 elsewhere), with both ontologies' modelling profiles in the prompt:

[system] You are OA-RLAIF, an ontology alignment agent. Decide which candidate target entity corresponds to the
source entity. … Reply with one line starting 'THOUGHT:' that names your evidence, then exactly one bash code block.
[user]   Task: find the target entity that the source entity is equivalent to.
         Alignment: anatomy/mouse-human
         ONTOLOGY PROFILES (modelling assumptions) …
         SOURCE
         IRI: http://mouse.owl#MA_0002211
         name: left renal vein
         parents: renal vein
         CANDIDATES
         [C1] … [C40] (each with name, synonyms, parents, children, definition)

An example reply (the format the mini-SWE-agent harness executes):

THOUGHT: C7 has the same name and the parent renal vein matches.
```mswea_bash_command
echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT && om submit C7 --relation = --confidence 0.95
```

Use thinking off (enable_thinking=False), as in training, and temperature 0.

The GRPO reward

For each prompt the policy samples G = 8 replies. Each reply gets

R = 0.8 · R_task + 0.2 · J / 10

Format gate. A reply must be one THOUGHT: line and exactly one om submit command whose choice is a valid candidate index (C1..Ck, at most 3) or NONE. Any other reply gets R = 0, and the judge is not called.

R_task, the task reward, depends on the kind of episode:

Episode kind	R_task
gold (the reference has the answer)	0.5 · F1(picks, gold) + 0.5 · r_DRAL
silver (ICD-11 links composed through WHO ICD-10, not gold)	the same, × 0.5
unlabeled (no reference)	0.3 · r_DRAL

F1 compares the submitted targets with the gold targets: one correct pick = 1; extra or missing picks lower it (recall counts at most 3 gold targets, the most one submission can name). NONE scores 1 only when no candidate is a gold target. If r_DRAL cannot be computed (a class without a stored vector), R_task = F1.

r_DRAL, the DRAL-OA reward (Eq. 1). For the source class s and each candidate k of the episode:

sim(s, k) = 0.5 · cos(z_s, z_k) + 0.5 · cos(t_s, t_k)
p_k       = softmax over the episode's candidates of sim(s, k) / T,   T = 0.1
r_DRAL(pick) = p_pick (mean over the picks)        r_DRAL(NONE) = 1 − max_k p_k

z is the DRAL-OA mapping-layer embedding of a class and t its label vector; both are stored in oa-rlaif-dral-reward-data. The mapping layer was fit on the training references only. r_DRAL gives graded credit: a wrong pick that DRAL-OA finds close to the source earns more than a far one.

J, the judge. google/gemma-4-26B-A4B-it, served locally next to the policy (thinking off), grades each valid reply from 0 to 10 without seeing the answer: evidence 0–4 (lexical, structural, semantic evidence that supports the choice), consistency 0–3 (the choice follows from the THOUGHT), faithfulness 0–2 (nothing invented), calibration 0–1 (the confidence fits the evidence). Rubric: configs/judge_rubric.md in the code repo.

Update. Dr. GRPO: advantage Â_i = R_i − mean(R) over the 8 replies of the prompt (no division by the standard deviation), β = 0 (no KL term), truncated replies masked.

Limitations

• Trained and evaluated with thinking off. With thinking on, the model often writes its answer inside the thinking text; the code repo's oa_rlaif.eval.thinking.split recovers it.
• Decides among retrieved candidates: a correct target that retrieval misses cannot be chosen.
• The image and audio parts of Gemma 4 are present in the weights but were neither trained nor used. The Hub may list the base architecture as image-text-to-text; this model is used for text generation only.

Licence and use. Private research use only; do not make this repo public. The weights follow the Gemma terms of use. The training data contains or derives from licensed sources: the OM datasource (Springer © 2007), SNOMED CT (affiliate licence) and UMLS-derived Bio-ML references.