s-ttp/aiva-rlcd-modernbert

AiVa RLCD: a request classifier for policy-constrained LLM routing

RLCD is the routing classifier in AiVa (AI Value Assurance). AiVa is a proof-of-concept enterprise AI gateway. Every request to the gateway first passes policy, data-loss-prevention (DLP) and quota checks. RLCD then reads the conversation and makes four predictions about it. The gateway uses them to choose the lowest-cost execution plan that is good enough for the request, but only among the models that policy already allows. RLCD can tighten a request's data classification. It can never loosen policy, budgets or routing constraints.

The model is a full fine-tune of answerdotai/ModernBERT-base. It has four linear classification heads on one shared, mean-pooled encoder, and each head is calibrated with its own temperature.

All training and evaluation data is synthetic. It describes GulfConnect Telecom (GCT), a fictional operator, and other invented organisations. The labels are produced by code from templates, so the accuracy below overstates what to expect on real traffic. See Evaluation and Limitations.

Model type	Encoder-only text classifier with four heads
Base model	answerdotai/ModernBERT-base at revision 8949b909ec90
Parameters	about 149 M (the encoder) plus four linear heads
Language	English
Maximum input	1,024 tokens (longer inputs are truncated)
Weights	fp32 safetensors, 596 MB
Licence	Apache-2.0
Developed by	The AiVa proof-of-concept team (s-ttp)

What it predicts

Head	Classes	What AiVa uses it for
task_category	qa, summarization, coding, rag, analysis, generation, agentic	The kind of work requested, for example to send coding requests to a capable model and to skip knowledge-base retrieval for them
sensitivity	public, internal, confidential, restricted, sovereign-only	A data-classification signal. It is used only to tighten the request's classification, never to relax it
complexity	simple, moderate, complex, advanced	How capable a model the request needs
context_need	sufficient, needs-retrieval, needs-packaging	Whether the prompt is self-contained, needs knowledge-base retrieval, or needs its context packaged by AiVa's context model (s-ttp/aiva-cim-qwen3-4b-v7)

Each head returns a calibrated probability distribution over its classes: softmax(logits / temperature), with the temperatures stored in rlcd_config.json.

How AiVa uses it

• Input. The conversation is rendered as one line per message, [role] content, in order and joined by newlines. Use the same rendering; the model was trained only on that format.
• Runtime. The model runs inside the gateway process on CPU (fp32, PyTorch). The gateway loads it once at startup, at a pinned revision of this repository.
• Policy comes first. RLCD chooses only within the set of models and paths that policy has already made eligible. Hard constraints are never learned and cannot be overridden by the model. Examples are sovereignty rules, model allowlists and budgets.
• Sensitivity only tightens. The gateway takes the highest level L for which P(class ≥ L) ≥ 0.5. The effective classification is then the most restrictive of four inputs: the declared level, the application's floor, DLP, and this prediction.
• Fallback. The gateway falls back to a rule-based router, and labels the run as such, in three cases: the model is unavailable, the model takes longer than its 2-second budget, or scoring fails. Because RLCD cannot weaken policy, falling back is always safe.

How to use

The repository holds the encoder configuration (encoder/), the tokenizer, the weights of the encoder and the heads (model.safetensors), and rlcd_config.json. That file lists the heads, class order, temperatures, pooling, maximum length and the measured test report. The checkpoint is not a stock AutoModelForSequenceClassification, so load it as below. This is the same loading and scoring the gateway uses.

# pip install "torch>=2.4" "transformers>=4.48" safetensors huggingface_hub
import json
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer


class RlcdClassifier(torch.nn.Module):
    def __init__(self, path: Path) -> None:
        super().__init__()
        self.config = json.loads((path / "rlcd_config.json").read_text())
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(path / "encoder"))
        hidden = self.encoder.config.hidden_size
        self.heads = torch.nn.ModuleDict(
            {name: torch.nn.Linear(hidden, len(classes)) for name, classes in self.config["heads"].items()}
        )
        # Parameter names match the trainer ("encoder.*", "heads.<name>.*"); strict on purpose.
        self.load_state_dict(load_file(str(path / "model.safetensors")), strict=True)
        self.eval()

    @torch.no_grad()
    def predict(self, messages: list[dict[str, str]]) -> dict[str, dict[str, float]]:
        text = "\n".join(f"[{m['role']}] {m['content']}" for m in messages)
        batch = self.tokenizer(text, truncation=True, max_length=self.config["max_length"], return_tensors="pt")
        hidden = self.encoder(**batch).last_hidden_state
        mask = batch["attention_mask"].unsqueeze(-1).to(hidden.dtype)
        pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1.0)  # masked mean pooling
        return {
            name: dict(zip(self.config["heads"][name],
                           (head(pooled) / self.config["temperatures"][name]).softmax(-1)[0].tolist()))
            for name, head in self.heads.items()
        }


model = RlcdClassifier(Path(snapshot_download("s-ttp/aiva-rlcd-modernbert")))
probabilities = model.predict([
    {"role": "user", "content": "What does the Traveller roaming bundle include?"},
])
for head, distribution in probabilities.items():
    label = max(distribution, key=distribution.get)
    print(f"{head}: {label} ({distribution[label]:.3f})")

Output (measured on CPU at revision cd53974b9859):

task_category: qa (0.933)
sensitivity: public (0.899)
complexity: simple (0.993)
context_need: needs-retrieval (0.986)

For a reproducible download, pin the revision: snapshot_download("s-ttp/aiva-rlcd-modernbert", revision="cd53974b98590daaf6d53b8952f6f9883d98a1c6"). Inference was developed with torch 2.14.1 and transformers 5.19.0.

Training data

The dataset is s-ttp/aiva-rlcd-labels at revision cd626dfed546. It is not published. It has 21,954 train, 2,720 validation and 3,727 test rows.

• Labels are fixed by code, not by people or a language model. They are derived by construction from the prompts in AiVa's two other synthetic datasets (the GCT domain dataset and the context-model dataset), plus templated sensitivity examples.

• Prose generator. The prose of those prompts was written by openai/gpt-oss-120b. That model was chosen because it is not one of the models RLCD routes to. A generated text was kept only if automatic checks passed:

  • every required fact is present;
  • there are no invented numbers, identifiers or timestamps;
  • chronology is preserved.

  Rejected texts were dropped, never repaired by hand.

• Identifiers. Subscriber-style identifiers (MSISDN, IMSI, ICCID, IMEI, IBAN, national ID) follow real formats, so DLP recognisers detect them. They are drawn at random and come from no subscriber list.

• Splits. Splits are assigned per group (document, case, customer and so on). A share of question templates is reserved for the test split, so no text the model trains on reappears in evaluation.

Training-split label counts:

Head	Counts (train)
task_category	analysis 7,502 · qa 6,507 · rag 3,239 · generation 1,490 · agentic 1,309 · coding 960 · summarization 947
sensitivity	internal 11,574 · confidential 6,520 · public 3,201 · restricted 519 · sovereign-only 140
complexity	moderate 8,882 · complex 8,182 · simple 3,958 · advanced 932
context_need	sufficient 12,434 · needs-packaging 6,464 · needs-retrieval 3,056

Training procedure

• Method. A full fine-tune of every encoder parameter and the four heads. The loss is the sum of the per-head cross-entropies; a row with no label for a head is masked for that head.

• Hyperparameters.

  • Optimiser: AdamW, learning rate 5e-5, weight decay 0.01.
  • Schedule: cosine, with 6% warm-up.
  • Gradient clipping at 1.0.
  • Batch size 32, 3 epochs, maximum length 1,024 tokens.
  • bf16 autocast.

• Calibration. One temperature per head, fitted by minimising NLL on the validation split only:

  Head	Temperature
  task_category	1.189
  sensitivity	0.939
  complexity	1.209
  context_need	1.336

• Hardware. One NVIDIA A100 on Hugging Face Jobs (a100-large). Training took 662 s.

• Software. torch 2.14.1 and transformers 5.19.0.

Evaluation

All numbers below are measured on the held-out synthetic test split of 3,727 rows. ECE is the expected calibration error, before and after temperature scaling.

Head	Accuracy	Macro-F1	ECE before	ECE after
task_category	0.991	0.983	0.0073	0.0057
sensitivity	0.964	0.781	0.0323	0.0320
complexity	0.879	0.857	0.0626	0.0528
context_need	0.976	0.977	0.0198	0.0166

CPU latency was measured on 150 test rows (fp32, 4 threads, up to 1,024 tokens): p50 352 ms and p95 746 ms per request.

How to read these numbers:

• The labels are synthetic and template-derived, so these figures overstate accuracy on real traffic. The model has not been evaluated on real enterprise traffic.
• The test split has no sovereign-only rows. That class therefore scores 0 in the sensitivity macro-F1, which explains most of the gap between its accuracy (0.964) and its macro-F1 (0.781). The class is also rare in training (140 rows).
• complexity is the weakest head (0.879). Its labels are the most subjective of the four.

Intended use

• Research and demonstration of policy-constrained, cost-aware LLM routing. In AiVa, the model only chooses inside the set that policy already allows.
• A starting point for fine-tuning on your own labelled traffic, with the same heads or new ones.

Out-of-scope use

• Do not use it as a data-classification, DLP or access-control mechanism on its own. The sensitivity head is a signal. In AiVa it can only tighten a classification, after policy and DLP have run.
• Security, compliance, legal or other high-stakes decisions, and any decision about people.
• Languages other than English. Inputs over 1,024 tokens are truncated.
• Domains far from its synthetic enterprise and telecom training distribution.

Limitations and risks

• Synthetic, template-derived training data. Real requests are more varied, and accuracy on them will be lower than the figures above.

• It is confidently wrong on unfamiliar phrasing. The prompts below were written by hand and do not come from the training templates. The model is run with the code above at revision cd53974b9859, and each cell shows its top class and that class's probability (measured). The task_category head often names the wrong task, and with high confidence. Generic coding and drafting requests in particular are read as agentic, analysis or rag. context_need says needs-retrieval for all seven, including self-contained coding requests. Sensitivity and simple factual questions held up better.

  Prompt	task_category	sensitivity	complexity	context_need
  Write a Python function that validates an IBAN and add unit tests for it.	agentic (0.733)	internal (0.963)	complex (0.655)	needs-retrieval (0.986)
  Write a Python function that parses a CSV of cell-site alarms and returns the five sites with the most critical alarms.	analysis (0.384)	internal (1.000)	complex (0.784)	needs-retrieval (0.982)
  Write a SQL query that returns the ten customers with the most dropped calls last week from the cdr table.	rag (0.760)	internal (0.991)	moderate (0.437)	needs-retrieval (1.000)
  Summarise the attached incident report for the operations manager in five bullet points.	analysis (0.559)	internal (1.000)	complex (0.833)	needs-retrieval (0.998)
  What does the Traveller roaming bundle include?	qa (0.933)	public (0.899)	simple (0.993)	needs-retrieval (0.986)
  Customer with MSISDN +97455123456 disputes a charge of QAR 120. Check the account and draft a reply.	qa (0.998)	confidential (0.785)	moderate (0.551)	needs-retrieval (0.826)
  Draft a short welcome email for new enterprise customers.	agentic (0.976)	internal (0.997)	complex (0.953)	needs-retrieval (0.998)

  In AiVa, a task that the caller or scenario declares takes precedence over this head. The gateway also uses the head only above a confidence threshold, and falls back to rules below it. Do the same, or fine-tune on your own traffic before relying on it.

• Rare classes are weak. sovereign-only and restricted are rare in training and absent or rare in the test split. Treat their probabilities with caution.

• Untrained success-by-plan heads. These heads would predict the outcome of each model and context strategy. They need measured outcomes and are not part of this checkpoint.

• Calibration is synthetic too. It was fitted on synthetic validation data and may not hold on real traffic. Re-fit the temperatures on your own data before relying on the probabilities.

Provenance and reproducibility

• Pinned inputs.
  • Base model: answerdotai/ModernBERT-base@8949b909ec900327062f0ebf497f51aef5e6f0c8.
  • Dataset: s-ttp/aiva-rlcd-labels@cd626dfed546c59f1af77242c5b3a39afd805896.
• rlcd_config.json records:
  • the architecture, pooling, heads and class order;
  • the temperatures and hyperparameters;
  • the test report, with each figure labelled measured.
• Weights revision. AiVa serves this repository at cd53974b98590daaf6d53b8952f6f9883d98a1c6. Later commits change only this card.
• Source code. The training and serving code is part of the AiVa proof of concept and is not public.