Reliable predictions in the wild: A confidence-aware prior injection for mitigating distributional risks

Few-shot classifiers deployed in safety-critical environments are highly vulnerable to distributional shifts, particularly when facing simultaneous label shift and open-set contamination. These conditions compromise decision integrity, causing models to misinterpret potential outliers or open-set samples as high-confidence known classes. A conventional defense is to estimate the test-time prior and inject it globally into classifier logits. However, we identify that this uniform correction introduces a brittleness in the inference pipeline: it creates a trade-off where strengthening the prior to correct in-distribution shifts inadvertently over-corrects ambiguous or adversarial-like queries, increasing the risk of harmful prediction flips. To harden few-shot systems against these distributional risks, we introduce Piecewise-π, a lightweight inference-time defense module. Unlike rigid global injection, Piecewise-π employs a confidence-aware mechanism that partitions queries based on their reliability. It applies a defensive, conservative correction to low-confidence queries (which are statistically likely to be outliers) while enforcing stronger calibration only on high-confidence, verified predictions. Experiments on CIFAR and STL benchmarks under severe imbalance and contamination demonstrate that Piecewise-π significantly improves the risk–coverage trade-off and system reliability without retraining the backbone. Behavioral analysis confirms that Piecewise-π effectively minimizes harmful decision reversals, offering a computationally efficient strategy to secure classifier performance in dynamic, open environments.