RAG-Gate-9B

RAG-Gate-9B sits after retrieval and before generation in a RAG pipeline. It reads a question, the retrieved passages, and whether more retrieval is possible, and emits one token: Answer (the evidence contains a complete support chain), Retrieve (it does not, and you can search again), or Stop (it does not, and you cannot). It is a LoRA fine-tune of Qwen/Qwen3.5-9B, merged into bf16 weights.

On a held-out test set of 14,818 items (2,256 distinct multi-hop questions), accuracy rises from .686 (same base model, same prompt, zero-shot) to .955. Most of the gain comes from the base model answering when the evidence is insufficient (.316 unsupported-answer rate); the fine-tune brings that down to .033 while refusing answerable questions only .062 of the time.

How to use

The decision is the first generated token after the prefill Final action:. Read the probabilities of the three label tokens directly; do not sample.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "ThakiCloud/RAG-Gate-9B"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16, device_map="auto")

POLICY = ("Policy: answer only if the retrieved evidence above contains a complete support chain for the "
          "answer. Do not use prior knowledge when judging whether the evidence is sufficient. "
          "If the evidence is insufficient and retrieval is available, retrieve more. "
          "If the evidence is insufficient and retrieval is not available, stop without answering.")
ACTIONS = "Actions: Answer = answer now; Retrieve = retrieve more evidence; Stop = stop without answering."

def gate(question, passages, retrieval_available=True):
    ev = "\n\n".join(f"[{i}] {p['title']}\n{p['text']}" for i, p in enumerate(passages, 1))
    user = (f"Question: {question}\n\nRetrieved evidence:\n{ev}\n\n"
            f"Retrieval available: {'YES' if retrieval_available else 'NO'}\n\n{POLICY}\n{ACTIONS}\n"
            "Reply with the action word only, on one line of the form 'Final action: <action word>'.")
    text = tok.apply_chat_template([{"role": "user", "content": user}], tokenize=False,
                                   add_generation_prompt=True, enable_thinking=False) + "Final action:"
    ids = tok(text, return_tensors="pt", add_special_tokens=False).to(model.device)
    labels = [tok.encode(w, add_special_tokens=False)[0] for w in (" Answer", " Retrieve", " Stop")]
    with torch.no_grad():
        logits = model(**ids).logits[0, -1, labels].float()
    p = torch.softmax(logits, -1).tolist()
    return dict(zip(("Answer", "Retrieve", "Stop"), p))

print(gate("Who directed the film that won Best Picture in 1998?",
           [{"title": "Titanic (1997 film)", "text": "Titanic won Best Picture at the 70th Academy Awards in 1998."}]))

Use Answer to let your generator write; Retrieve to run another retrieval round; Stop to return "I can't answer from the available documents". You can threshold p["Answer"] instead of taking the argmax if your application prefers fewer unsupported answers over more refusals.

What changes — real test-set examples

Each row is a test question where the base model chose wrong and RAG-Gate-9B chose right (picked deterministically by item-id hash; passages omitted for space).

One error, also picked by hash: "Who is the programming language that has the WHERE clause partially named after?" — state FULL, retrieval YES; the correct action is Answer, RAG-Gate-9B said Retrieve.

Results (blind test, 14,818 items over 2,256 base questions; 95% CI by bootstrap over base questions, 10,000 resamples)

By evidence state (accuracy):

FULL_DECOY matters most: the evidence was edited but is still sufficient, so the right action is Answer. A model that learned "edited text means refuse" would fail here.

ChainCheck (out-of-distribution, built separately): pairs that test whether the model reacts to whether the support chain is intact (CE) more than to surface edits (EE). Σ = CE − |EE| should be positive.

All numbers were measured by us, with the prompt above and bf16 weights, on our own GPUs. We do not compare against other vendors' models here.

Release gates (pre-registered before training)

The model was released only because it passed all five gates, fixed before training started:

G2 was originally "unsupported rate below zero-shot". We replaced it before full training: in the smoke test the 4B zero-shot model refused almost everything, so nothing could beat it on that metric and a model that always refuses would win. For this model unsupported answers also fell, from .316 to .033.

Limitations

• English only, one source domain. Training and test data are derived from MuSiQue (Wikipedia, 2–4 hop questions). Korean, enterprise documents, tables, and code have not been measured.
• The test set is in-domain. Blind test shares the construction procedure with training (different base questions). ChainCheck is the only out-of-distribution check.
• It judges sufficiency, not truth. It is told not to use prior knowledge; a passage that is wrong but internally complete is judged sufficient.
• Text-only. The base model is multimodal; the vision tower was not trained and is not included.
• Long inputs. Inputs longer than 2,048 tokens were not evaluated (8 of 14,818 test items were dropped for length).
• Merging into bf16 changes probabilities slightly (max |Δp| 0.046 on the G5 sample); decisions were unchanged on that sample.

Training

LoRA r=16, α=32, all linear layers; loss on the single label token only; 32,768 training rows (sampled by base question), 1023 steps, effective batch 32, lr 0.0001, linear warmup/decay, max length 2048. Checkpoint selected on a separate calibration slice (step 1023). 1× GPU.

Data

Built from MuSiQue (CC BY 4.0) by deleting, contradicting, or editing passages to create the five evidence states, crossed with the retrieval-available bit. No personal data and no AI Hub data are included. The training data is not distributed with this model.

Related

• ChainCheck — the counterfactual benchmark used for G4.
• ChainCheck-Judge — a scalar sufficiency score (log-odds) instead of a three-way action.

License

Apache-2.0, same as the base model.