Trust-Gated State Space Models for Budgeted Decisions with Online Risk Control

Kailash Talreja, Saurabh Jha

Keywords: reliable machine learning, decision making, budgeted abstention, conformal prediction, uncertainty calibration, online learning, distribution shift, state space models, Mamba, risk control, edge deployment

TL;DR: A tiny SSM with reliability gates plus online, recency-weighted conformal prediction that optimizes budgeted decisions while keeping miscoverage near target under drift.

Abstract: Deployed predictors increasingly operate with unreliable inputs and strict resource budgets, where labels can be noisy, sensors fail, and conditions drift. We introduce Trust-Gated State Space Models (TG–SSM), a compact approach that treats reliability as a first-class control signal. TG–SSM augments a lightweight statespace backbone with gates that modulate input injection, state mixing, and output temperature using on-the-fly reliability features; a small conformal layer then converts probabilities into calibrated prediction sets for budgeted decision-making with target risk 1 − α. On CIFAR-10N (noisy labels), TG–SSM with weighted conformal prediction (WCP) reduces ECE from 0.124 to 0.048 while increasing coverage from 0.894 to 0.905 at α = 0.1 (mean set size ≈ 9.08). Averaged over CIFAR 10C severities 1–5, TG–SSM+WCP achieves near-nominal coverage (0.905) with markedly improved calibration (ECE 0.104–0.109) and compact sets (≈ 6.74). On Camelyon17/WILDS (domain shift), a validation-quantile variant attains AUROC 0.949 with average coverage 0.844 and set size ≈ 0.989, while a shift-aware (importance-weighted) variant yields smaller sets (≈ 0.737) with AUROC 0.936. Overall, TG–SSM provides a simple, hardware-efficient recipe for turning uncertain predictions into actionable, budgeted decisions.

Primary Area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)

Submission Number: 23100