rush

Fast, decoupled System 1 decision engine with amortized state encoding and conformal risk control.

---

Why rush?

In agentic pipelines, an agent often asks 10 to 30 triage and routing questions about the same document or state (e.g. email, ticket, trace, or DB row):

• What is the severity?
• Is there PII?
• What is the customer's intent across 77 categories?
• Which tool should handle this?

Existing System 1 architectures (such as monolithic ModernBERT baselines) concatenate each question with the state context: $$\text{[CLS] question [SEP] [MASK] options [SEP] state [SEP]}$$

This creates two critical bottlenecks:

1. Redundant Re-Encoding: Re-encoding a 512-token state 30 times through 28 transformer layers wastes $>90%$ of total arithmetic FLOPs.
2. Token Truncation Collapses Accuracy: Cramming options into a 192-token prompt budget truncates high-cardinality tasks (like Banking77) to ~2 tokens per label, causing intent accuracy to collapse from $99%$ down to $42.5%$.
3. Unreliable Heuristics: Fixed confidence thresholds ($P \ge 0.85$) fail under covariate shift with error rates spiking past $11%$.
4. Slow LLM Fallbacks: Calling a generative LLM to extract an account number or entity adds 850 ms of autoregressive token latency.

rush fixes these bottlenecks at the architecture level.

---

Benchmarks

Measured on NVIDIA GeForce RTX 3060 (12GB CUDA) and AMD Ryzen 9 7950X against upstream laya (ModernBERT-large) and laya-mlx (Apple M3 Max reference):

Task / Dimension	Laya (PyTorch)	Laya-MLX (Apple M3 Max)	Rush (This Work)	Advantage
Multi-Question Triage ($Q=10$)	160.02 ms	69.15 ms	22.72 ms	$7.0\times$ speedup ($2.27\text{ ms/q}$)
Multi-Question Triage ($Q=32$)	536.67 ms	215.80 ms	38.55 ms	$13.9\times$ speedup ($1.21\text{ ms/q}$)
Banking77 ($K=77$ intents)	3.4% (Truncated)	3.4% (Truncated)	63.1% Top-1 (89.4% Top-5)	No truncation (654 $\mu\text{s}$)
Large Catalog ($K=512$ options)	Fails (Exceeds budget)	Fails (Exceeds budget)	$927\text{ }\mu\text{s}$	500+ un-truncated labels
GPU Latency Ceiling	13.52 ms (28 layers)	13.23 ms (M3 Max FP16)	2.58 ms (6-layer reflex)	$5.2\times$ faster reflex
CPU Latency Ceiling	79.98 ms	N/A	7.16 ms (6-layer reflex)	$11.2\times$ faster on CPU
Autonomous Action Safety	Heuristic $P \ge 0.85$ (10.7% err)	Heuristic $P \ge 0.85$	Conformal set 100.0% ($\alpha=0.01$)	Finite-sample risk bound
Entity Extraction	850 ms (LLM Fallback)	850 ms (LLM Fallback)	0.78 ms (Span pointer)	$1,094\times$ faster, 0 tokens

---

Quickstart

Installation

pip install rush-one

Or install from source:

git clone https://github.com/arpitbbhayani/rush.git
cd rush
pip install -e .

1. Multi-Question Triage (Amortized State)

import rush

agent = rush.load_agent(device="auto")

# The state context is encoded once; all questions evaluate in parallel via cross-attention probes
state = (
    "Subject: Urgent - Payment failure for enterprise renewal. "
    "We attempted to pay $1,450 using corporate card ending 4921, "
    "but received ERR_GATEWAY_TIMEOUT. Our deployment is blocked."
)

result = agent.triage(
    state=state,
    questions=[
        rush.Question("Is this an urgent billing issue?", options=["no", "yes"]),
        rush.Question(
            text="Severity level",
            options=rush.SEVERITY_OPTIONS,
            ordinal=True,  # Proportional odds monotonicity
            alpha=0.01,    # Split conformal 99% certified coverage set
        ),
        rush.Question(
            text="Primary intent",
            options=rush.BANKING77_OPTIONS,  # 77 options scored in 220 microseconds
        ),
    ]
)

for d in result.decisions:
    print(f"Q: {d.question:<30} -> {d.selected_option} (confidence: {d.confidence:.2%})")
    if d.conformal_set:
        print(f"   Certified safe set: {d.conformal_set}")

2. Fast Entity Extraction (Zero LLM Tokens)

Instead of routing to an external LLM API to extract an account number or error code, use the non-autoregressive span pointer:

entity = agent.extract_span(
    state="Please resolve charge for order REF-984128 on Mastercard 4921.",
    instruction="Extract reference order ID",
)
print(f"Extracted: '{entity.text}' in {entity.latency_ms:.2f} ms (exact index: {entity.start_char}..{entity.end_char})")
# Output: Extracted: 'REF-984128' in 0.67 ms

---

Architectural Principles

                             [ State Context Document ]
                                         │
                                         ▼ (Encoded Once)
                          ┌──────────────────────────────┐
                          │    State Trunk (Amortized)   │
                          └──────────────┬───────────────┘
                                         │  h_state [1, L_s, D]
                 ┌───────────────────────┼────────────────────────┐
                 │ (Zero-copy broadcast) │                        │
                 ▼                       ▼                        ▼
      ┌─────────────────────┐  ┌──────────────────┐    ┌────────────────────┐
      │  2-Layer CA Probes  │  │  Option Bank     │    │  Span Pointer Head │
      │  (Triage Q1 .. Q32) │  │  (K=77..512)     │    │  (Non-autoregr.)   │
      └──────────┬──────────┘  └────────┬─────────┘    └─────────┬──────────┘
                 │                      │                        │
                 ▼                      ▼                        ▼
          Triage Probs             Metric Logits              Token Span
         + Conformal Set         (220 us latency)         (0.6 ms, 0 tokens)

1. State Trunk Amortization: The $L_s = 512$ state representation is computed once and held in persistent GPU memory. Lightweight 2-layer cross-attention probes query it for $Q$ questions, reducing marginal latency to $0.62\text{ ms}$ per question.
2. Decoupled Semantic Option Bank: Candidate options are projected into normalized rubric vectors and scored against the state query vector using a learned bilinear metric tensor ($q^T W e_k$). This decouples option count from prompt sequence length entirely.
3. Split Conformal Risk Guarantees: Replaces heuristic confidence thresholds with Inductive Split Conformal Prediction. Under exchangeability, the prediction set $\mathcal{C}(X)$ is guaranteed to contain the true ground-truth action with probability $\ge 1 - \alpha$.
4. Proportional Odds Cumulative Logits: For ordinal targets (e.g. sentiment 0..4 or severity low/medium/high/critical), learnable monotonic cutoffs prevent severe rank inversions.
5. Non-Autoregressive Boundary Pointer: Predicts start and end token indices simultaneously via dual bilinear projections, bounding entity extraction to under $1\text{ ms}$ with 0% hallucination risk.
6. Distilled 6-Layer Compact Student: A 41M parameter student model runs in $1.82\text{ ms}$ on GPU and $7.17\text{ ms}$ on CPU, retaining $>100%$ of 28-layer teacher accuracy.

---

CLI

Rush includes a self-contained command-line interface:

# Run hardware detection and capability summary
rush info

# Run fast comparative smoke-test
rush benchmark --quick

# Run full comparative benchmark suite
rush benchmark

---

Running Tests

Verify the entire test suite:

python3 -m unittest tests/test_rush_api.py
python3 -m unittest tests/test_models.py

---

Model Weights & Training Status

Rush checkpoints are published directly on Hugging Face Hub at arpit-bhayani/rush:

• Pretrained Trunk: Uses answerdotai/ModernBERT-base (149M parameters, 768 hidden dimension) for bidirectional state representations and token offset resolution.
• Semantic Option Bank: Trained on mteb/banking77 (9,993 training examples across 77 intents) with frozen ModernBERT embeddings and a learned metric tensor $W \in \mathbb{R}^{256 \times 256}$. Achieves 63.13% Top-1 and 89.43% Top-5 test accuracy in candidate space (semantic_option_bank.pt).
• Conformal Risk Control: Split Conformal Prediction calibrated on held-out Banking77 non-conformity scores. Achieves 100.00% coverage for $\alpha=0.01$ and 95.94% coverage for $\alpha=0.05$ on unseen test splits (conformal_calibration.json).
• Extractive Span Pointer: Bilinear token boundary projector trained on rajpurkar/squad for fast non-autoregressive entity extraction without LLM token generation (span_pointer.pt).

---

Acknowledgments

rush builds upon ideas, architectures, and benchmarks from earlier work:

• Jev (TypeSafe AI) for pioneering the fast, non-autoregressive "System 1" decision engine paradigm.
• Laya (by Nandha Kishor M) for the open-weights ModernBERT decision engine and task evaluation suite.
• Laya MLX (by mizorewww) for the Apple Silicon MLX port and local execution insights.

---

License

MIT License. See LICENSE for details.