budgeted-exit-qwen3.5-2b-civil-comments

Budgeted early exit for Qwen3.5-2B (Civil Comments moderation)

Stop computing when the decision is settled. This is a LoRA fine-tune of the hybrid linear-attention model Qwen3.5-2B with a trained classification head at seven exit layers and a frozen, risk-calibrated exit policy. It answers four moderation questions about a comment (threat, identity attack, insult, obscene language), composes them with a rule, and leaves the network early when the rule's outcome can no longer change within an explicit error budget.

On a locked, pre-registered test it keeps decision quality within 0.2 points of the full model while running 5.45 of 24 layers on average, and a runtime that really stops is 3.99x faster in measured throughput.

Research release. Trained and evaluated on one dataset (Civil Comments) and one seed. Read the limitations before using it.

Results

Rule: threat or identity_attack or (insult and obscene) (about 1.2% positive). Official Civil Comments test split, 96,862 comments, read once after the policy, weights, code and test file were hashed. Pass rule fixed in advance: the 95% interval of the change in balanced accuracy must exclude a 1-point loss.

Balanced accuracy of each exit on the validation split, if every comment stopped there:

The policy changed 2.54% of the full model's positive decisions (budget 2%) and 0.92% of its negative decisions (budget 1%). The first figure overshoots its budget: the isotonic maps and the dials were fitted on the same calibration rows. Treat the budget as approximate.

How it works

1. Each exit head outputs a probability per question; Platt scaling maps it to the natural label distribution.
2. The rule's probability is computed from the four probabilities (weighted model count, independence assumed).
3. Per exit, an isotonic map gives the probability that the full-depth decision would be positive. It is fitted without labels.
4. One dial per direction, scaled by the share of layers still to run, decides who leaves. The dials are set by conformal risk control against the error budget. Everything is stored in frozen_policy.json as plain numbers.

Files

• model.safetensors: LoRA adapters (rank 16, all large linear maps) and the seven exit heads, 26.7M parameters. The base weights are downloaded from Qwen/Qwen3.5-2B.
• config.json: base model, exit layers, questions, LoRA rank, prompt prefix and suffix
• frozen_policy.json: the exact policy evaluated on the locked test (Platt parameters, isotonic maps, dials, cut-offs)

Usage

The loading and runtime code lives in the project repository: github.com/hamdallah90/decision-aware-inference (folder budgeted_exit/).

git clone https://github.com/hamdallah90/decision-aware-inference.git && cd decision-aware-inference
pip install -r requirements.txt && cd budgeted_exit

from huggingface_hub import snapshot_download
from hub import load                      # budgeted_exit/hub.py
from runtime_exits import Runtime         # budgeted_exit/runtime_exits.py

folder = snapshot_download("hamdallah/budgeted-exit-qwen3.5-2b-civil-comments")
model, encode, _ = load(folder)           # logits = model(ids, mask) -> [batch, exits, 4 questions]
runtime = Runtime(tag=None, folder=folder)
decisions, exit_layers = runtime.dynamic(encode(["your comment here"]))

Training

LoRA (rank 16) on a frozen bf16 backbone, last-token readout, one sequence per comment with all four questions in a fixed prompt prefix, on 55,271 Civil Comments training rows (an enriched sample with extra rare positives), binary cross-entropy summed over all exits, two epochs, one seed. Calibration used a separate uniform sample of 99,140 training rows.

Limitations

• One dataset, one seed, one hand-written rule. Re-measure on your own data before relying on it.
• Throughput was measured with requests available in bulk. Latency under live traffic was not measured.
• The error budget is controlled in expectation on data like the calibration sample; it is not a per-request guarantee and it does not cover distribution shift.
• English only. The training data contains offensive language; the model is for moderation research, not for generating content or for automated decisions about people without human review.
• The runtime calls Hugging Face layers directly and was tested with transformers 5.17.

License

Apache-2.0, following the base model. Dataset: Civil Comments (CC0-1.0).