Ajeet-001/embeddinggemma2-decision-router

embeddinggemma2-decision-router

embeddinggemma2-decision-router is a System-One multimodal decision and routing model built on google/embeddinggemma-2. Given a workflow state (text, image, or audio) and a schema of questions, it returns a typed, calibrated probability distribution over every option in a single forward pass:

• Zero option-order bias: 0.0% order-flip rate across option permutations (set-equivariant transformer head with no positional encoding on options).
• Calibrated Null Gate & low latency: 0.9163 out-of-scope abstention AUROC and 53.6 ms median latency on 1x A100 (bfloat16).

Question type	Example	Python type of value
choice	What should we do with this trace? → continue / observe / human_review / stop	str
score	How risky is it? → level 0..3	int
noul	Does it need human review? → probabilistic yes/no	bool

---

📋 Model Details

Model	Parameters	Description
google/embeddinggemma-2	270M (Text) / 740M (Multimodal)	Base multilingual & multimodal embedding backbone (frozen)
embeddinggemma2-decision-router	+3.55M Decision Head	2-layer Set-Equivariant Transformer + 6-feature Calibrated Null Gate

• Inputs: text, images, and 16 kHz audio (shared 768-d EmbeddingGemma 2 embedding space).
• Architecture: frozen EmbeddingGemma 2 encoder + Pre-LN Set-Equivariant Transformer (d_model=768, n_heads=8) + Platt-calibrated Null Gate.
• Multilingual: inherits EmbeddingGemma 2's 100+ language alignment because the encoder remains frozen.
• Precision: use torch.bfloat16 or torch.float32 (never float16).

---

🏃 How to use

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "Ajeet-001/embeddinggemma2-decision-router",
    trust_remote_code=True,
    dtype=torch.bfloat16,
).to("cuda")

state = {
    "agent": {"autonomy": "unsupervised", "model": "internal-agent-v2"},
    "task": "Clean up stale feature branches merged more than 90 days ago.",
    "constraints": ["Never delete anything matching release/*", "Do not force-push to any branch"],
    "trace_summary": {"steps": 9, "tool_errors": 1, "constraint_violations": 1, "irreversible_actions": 1},
}

questions = {
    "action": {
        "type": "choice",
        "instructions": "What should the observability system do with this trace?",
        "criteria": {
            "continue": "Let the agent proceed without interruption.",
            "observe": "Keep running, but flag the trace for later sampling.",
            "human_review": "Queue this trace for a human to review.",
            "stop": "Halt the agent now.",
        },
    },
    "needs_review": {"type": "noul", "instructions": "Does a human need to review this trace?"},
    "risk": {
        "type": "score",
        "instructions": "How risky is what the agent did?",
        "criteria": ["no risk", "low risk", "significant risk", "severe risk"],
    },
}

answers = model.decide(state, questions, workflow="agent_trace_observability")

for qid, answer in answers.items():
    print(f"{qid:13s} → {answer.value!r:15} ({answer.confidence:.0%})  {answer.probabilities}")

Every answer is a TypedAnswer:

TypedAnswer(
    kind="choice",                  # "choice" | "score" | "noul"
    label="observe",                # winning option key
    value="observe",                # str for choice, int for score, bool for noul
    confidence=0.5037,              # calibrated probability of the winning option
    probabilities={"continue": 0.1908, "observe": 0.5037, "human_review": 0.2771, "stop": 0.0284},
    expected_level=None,            # score questions only: E[level]
)

---

📊 Performance

Jev Decision Index (Public Split)

Suite	Score	Median Latency
Vision Index 0.3 (12,210 rows)	7.80	130.8 ms
Text Index 0.3 (140,178 rows)	6.99	53.6 ms

Zero-shot highlights on the public split: BANKING77 (76.61%), Winoground (64.25%), CV-Bench (53.88%), Hateful Memes (53.01%), CLINC150+OOS (50.35%), SGD/SGD-X (48.59%), and ARC-Easy (48.19%). Full run artifacts: Ajeet-001/embeddinggemma2-decision-index-results.

Controlled Candidate Routing (BANKING77, K=8 Hard Negatives)

Model	Seen Acc	Unseen Zero-Shot	Null AUROC	Brier
Zero-Shot Cosine + Platt Gate	84.33%	89.56%	0.7793	0.0365
embeddinggemma2-decision-router	90.35%	85.44%	0.9163	0.0258

Zero-Shot Multimodal Routing (Vision & Audio)

Modality	Benchmark	Score
Vision (740M tower)	RVL-CDIP (Document Routing, K=4)	83.00%
Audio (740M tower)	SUPERB KS (Spoken Command Routing, K=8)	82.50%

---

🔧 Training

• Data: 15,884 K=8 hard-negative routing episodes (Ajeet-001/banking77-hard-negative-routing-episodes). 70% covered multi-choice episodes + 30% out-of-scope abstention episodes built from the 62 seen classes of PolyAI/banking77 (train split; 15 classes held out zero-shot).
• Objective: cross-entropy over candidates + binary cross-entropy on the decoupled Null Gate + Platt temperature calibration.
• Recipe: 15 epochs, peak LR 3e-4, AdamW (weight_decay=1e-2), cosine decay, bf16, encoder 100% frozen.
• Hardware: 1 × NVIDIA A100-SXM4-40GB.