PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Nonlinearities Improve Uncertainty Estimation

Mohammad Sarhangzadeh, Amirreza M. Shebly, Doruk Oner

arXiv:2610.29875v1cs.CV

TL;DR

Visual models can be overconfident on ambiguous or out-of-distribution inputs, while uncertainty methods may suppress features or require substantial computation. This paper estimates uncertainty by sampling activation functions in one model, and reports improved uncertainty quality with competitive task performance across visual settings.

  • Problem

    Visual models can make overconfident predictions on ambiguous, corrupted, or out-of-distribution inputs, while uncertainty methods face trade-offs between computational efficiency and preserving learned representations.

    Source: §1 · pp. 1–2

  • Method

    A single trained model samples between nonlinear activation functions during repeated forward passes, using predictive disagreement to estimate uncertainty while keeping feature pathways active.

    Source: §1 · pp. 2–3

  • Results

    Across classification, segmentation, and out-of-distribution settings, stochastic activations improve uncertainty quality while maintaining competitive task performance; MiniImageNet classification accuracy is 66.90 vs. 70.50 for Deep Ensembles.

    Source: §4.2 · p. 11, §6 · p. 14

  • Takeaways & Limitations

    Stochastic nonlinearities provide a lightweight, information-preserving source of predictive variation from one trained model without separate uncertainty heads or task-specific architectural redesign.

    Source: §6 · p. 14

Abstract

from arXiv · show

Reliable uncertainty estimation is important across visual recognition tasks, providing information about the confidence and reliability of model predictions. Many uncertainty estimation methods obtain predictive diversity through mechanisms that can suppress intermediate features, perturb the input, or require substantial additional computation. Stochastic activation functions, where network units randomly select between multiple nonlinearities during a forward pass, have previously been used to improve optimization behavior, mitigate limitations of fixed ReLU activations, and increase output diversity. In this work, we reframe stochastic activation switching as a lightweight mechanism for uncertainty estimation. The proposed method samples between nonlinear activation functions, inducing predictive diversity while keeping intermediate feature pathways active. This enables uncertainty estimation from a single trained model without the training cost of ensembles, task-specific uncertainty heads, or architectural redesign. We evaluate the proposed approach across classification, semantic segmentation, and out-of-distribution detection using convolutional and transformer-based architectures on diverse datasets. Across these settings, stochastic activation switching improves uncertainty quality while preserving competitive primary-task performance, showing that stochastic nonlinearities provide an effective and information-preserving source of predictive variation for visual recognition.

1 Introduction

Reliable uncertainty estimates matter because overconfident visual errors can undermine downstream decisions, while existing approaches trade computational cost against preserving learned features. This work uses stochastic activation switching to estimate uncertainty from one model and evaluates it across visual tasks and architectures.

1 Introduction

  • Visual models can be overconfident on ambiguous, corrupted, or out-of-distribution inputs, making unreliable predictions a concern for downstream recognition pipelines.Uncertainty estimates augment predictions with confidence measures that help identify unreliable outputs.Source: §1 · p. 1
  • Ensembles require training and storing multiple models, while dropout-based diversity suppresses features and can disrupt representations needed for classification or dense prediction.These trade-offs motivate diversity mechanisms that preserve learned feature representations.Source: §1 · p. 2
  • Stochastic activation methods randomly select among nonlinear functions during forward passes, but prior work primarily studied them for optimization and regularization rather than uncertainty estimation.The paper reframes them as a way to introduce predictive variation while preserving information flow through the network.Source: §1 · p. 2
  • A single trained model uses repeated stochastic activation passes to produce predictive disagreement, without multiple-model training or substantial architectural redesign.The framework is evaluated with convolutional and transformer-based vision models.Source: §1 · p. 3
  • Experiments span classification and segmentation across diverse datasets, modalities, scales, and in- and out-of-distribution settings, with competitive task performance and strong uncertainty estimates.The paper also reports ablations on activation and inference settings.Source: §1 · p. 3

2 Related Work

Existing uncertainty methods obtain prediction diversity through ensembles, feature or weight masking, deterministic feature geometry, or other inference-time perturbations. Prior stochastic-activation work focused mainly on training behavior; this paper applies dynamic activation sampling within one model to uncertainty estimation.

2.1 Uncertainty Estimation

  • Deep ensembles estimate uncertainty from independently trained networks but incur training, storage, and inference costs that grow with ensemble size.MC Dropout is a lighter single-model alternative, but generates diversity by suppressing features; related variants also mask weights or learn dropout rates.Source: §2.1 · pp. 3–4
  • Deterministic methods such as DUQ, SNGP, and DDU avoid repeated stochastic inference, while other approaches use input transformations, normalization noise, iterative refinement, or geometric constraints.These alternatives depend on learned feature geometry or method-specific design choices.Source: §2.1 · p. 4

2.2 Stochastic activations

  • Stochastic activations have been used for regularization and optimization, including addressing dying ReLUs, and later work explored training dynamics, sparsity, runtime adaptability, and generative diversity.Lomeli et al. [38] used test-time stochasticity to increase language and generative output diversity.Source: §2.2 · p. 4
  • Random Activation Functions estimates uncertainty across independently trained models with fixed activation choices, whereas this method samples activations dynamically within one model during training and inference.Its uncertainty comes from disagreement across stochastic passes and is not decomposed into epistemic and aleatoric components.Source: §2.2 · p. 4

Method

The method replaces fixed activations with stochastic switching to create prediction diversity without removing features. It uses prediction variability from the same trained model to estimate uncertainty.

3.1 Motivation

  • Stochastic activation switching transforms signals instead of removing features, preserving active pathways without uncertainty heads, extra learnable parameters, or repeated model training.The approach targets diverse predictions from one trained model without duplicating the network or adding trainable parameters.Source: §3.1 · p. 5

3.2 Stochastic Activation Switching

  • Each unit switches between the backbone activation and an alternative differentiable nonlinearity with probability p, which defaults to 0.5 for uniform switching.The backbone activation may be ReLU in convolutional networks or GELU in transformer-family models.Source: §3.2 · p. 5
  • Sampling can occur per activation unit, per channel shared across spatial locations, or once for an entire layer.These node-, channel-, and layer-level choices provide different forms of functional diversity.Source: §3.2 · p. 6
  • In toy heteroscedastic regression, trajectories stay stable in supported regions and vary more in noisier regions and outside the training range than MC-Dropout.Both methods use repeated stochastic passes in the same setup.Source: §3.2 · p. 7, Figure 3 · p. 7

3.3 Training and Uncertainty Estimation

  • Training and inference use the same switching mechanism; fixed weights then generate T stochastic predictions whose empirical variance estimates class-probability or pixelwise uncertainty.For segmentation, the variance produces a spatial uncertainty map.Source: §3.3 · pp. 7–8
  • Information Destruction Rate (IDR) compares layerwise input–output mutual information under deterministic and stochastic inference; lower IDR indicates greater information preservation.IDR is a diagnostic comparison, not a training objective, including against feature-masking baselines such as MC Dropout.Source: §3.3 · p. 8

4 Experiments

The experiments test stochastic activation switching across segmentation and classification datasets, architectures, and distribution shifts. Results show competitive task performance alongside strong uncertainty estimates, with MiniImageNet classification accuracy as a noted exception.

4.1 Experimental Setup

  • Experiments use segmentation and classification datasets, compare common uncertainty baselines, and evaluate task performance, uncertainty quality, and OOD behavior under matched within-task protocols.The study includes U-Net-based and Swin Transformer V2 segmentation models, ResNet-style and Swin Transformer V2 classifiers, and task-specific metric suites.Source: §4.1 · pp. 8–9
  • For segmentation visualization, the uncertainty map is computed from variance across five stochastic activation trajectories.The figure presents input image, ground-truth mask, prediction, and uncertainty map from left to right.Source: Figure 4 · p. 9

4.2 Results

  • Across segmentation benchmarks, stochastic activation switching ranks best or second-best on most task and uncertainty metrics while retaining the original architecture and single-model training setup.The results cover thin structures and biomedical settings, with uncertainty aligning well with prediction errors and strong calibration and Brier scores.Source: §4.2 · pp. 9–10, Table 1 · p. 10, Table 2 · p. 11
  • Under cross-dataset segmentation shifts, disagreement between stochastic activation trajectories remains informative.Models trained on Massachusetts, DRIVE, and CREMI are tested on RoadTracer, CHASE-DB1, and ISBI 2012, respectively.Source: §4.2 · p. 10, Table 3 · p. 12
  • Classification results show improved uncertainty behavior with strong task performance, although MiniImageNet accuracy is 66.90 versus 70.50 for Deep Ensembles despite better calibration and OOD metrics.On high-resolution medical benchmarks, the method transfers to Swin Transformer V2, maintaining high in-distribution accuracy while improving cross-domain OOD separation.Source: §4.2 · pp. 10–11, Table 4 · p. 12
  • Across most settings, stochastic activations improve uncertainty quality while maintaining competitive task performance across U-Net, ResNet-style, and Swin Transformer V2 architectures.The approach uses one trained model without an uncertainty head, architectural redesign, or multiple-model training.Source: §4.2 · p. 12

Ablation Study

Ablations examine stochasticity granularity, activation choice, information preservation, and inference-pass count. They motivate node-level sampling, two-activation switching, and five inference passes as practical design choices.

Granularity of Stochasticity

  • Node-level sampling provides the strongest overall trade-off across predictive performance and uncertainty quality, improving accuracy, NLL, and Brier score, while layer-level sampling gives the lowest ECE.The authors attribute richer functional diversity and reduced local activation co-adaptation to the more local stochasticity.Source: Ablation Study · p. 13, Table 5 · p. 13

Choice of Stochastic Activation

  • ReLU+SiLU gives the best accuracy and NLL, while ReLU+Tanh gives the best ECE and Brier; more than two candidate activations was less stable.SiLU is fixed for classification and Tanh for segmentation across datasets, modalities, and architectures; other pairs remain competitive.Source: Ablation Study · p. 13, Table 6 · p. 13

Information Destruction Analysis

  • Stochastic activation switching gives lower IDR than MC Dropout across datasets, indicating better preservation of intermediate representations under stochastic inference.The method perturbs nonlinear transformations while keeping feature pathways active.Source: Ablation Study · pp. 13–14, Figure 5 · p. 14

Number of inference passes

  • Increasing inference passes stabilizes predictive means and uncertainty estimates and improves uncertainty metrics, while increasing runtime approximately linearly and leaving accuracy saturated after a small number of samples.The experiments use T = 5 as a practical balance between uncertainty quality and computational cost.Source: Ablation Study · p. 14, Table 7 · p. 14

6 Conclusion

Stochastic activation switching improves uncertainty quality across classification, segmentation, and out-of-distribution detection while preserving main-task performance. Its calibration–performance trade-off depends on the activation pair, and repeated stochastic passes add inference cost.

6 Conclusion

  • Across classification, segmentation, and out-of-distribution detection, stochastic activation switching improves uncertainty quality without degrading main-task performance.The conclusion reports effectiveness across visual domains and architectures.Source: §6 · p. 14
  • Sampling between nonlinear activations keeps intermediate feature pathways active and provides predictive variation from a single trained model.The method avoids separate uncertainty heads and task-specific architectural redesign.Source: §6 · p. 14
  • The activation pair affects the trade-off between predictive performance and calibration, while repeated stochastic passes increase inference cost.The conclusion identifies adaptive activation selection and more efficient sampling as future directions.Source: §6 · p. 14

Supplementary Material

Stochastic Activation Switching replaces deterministic activations with per-unit Bernoulli choices between a backbone’s native activation and a fixed alternative. Training samples activation configurations while learning weights; inference samples repeated predictions with fixed weights to estimate a mean and variance.

A Approach Overview

  • Each activation unit switches between the backbone’s native activation and a fixed alternative using a Bernoulli choice during training and inference.The same switch changes the nonlinearity without removing features from intermediate representations.Source: A Approach Overview · p. 21
  • Training samples a fresh activation configuration on every forward pass and optimizes the network with the standard task loss.The learned weights are thereby trained against activation-level variation rather than one fixed nonlinearity.Source: A Approach Overview · p. 21
  • At inference, fixed trained weights produce T stochastic predictions per input; their empirical mean is the point estimate and their variance is the uncertainty estimate.The perturbation changes nonlinear transformations while keeping all activation pathways active, distinguishing it qualitatively from dropout-style feature removal.Source: A Approach Overview · p. 21

B Implementation Details

Implementation uses task-level activation choices and node-level Bernoulli switching, with training procedures matched to deterministic backbones. Experiments span multiple architectures, datasets, and baseline configurations; inference uses five stochastic passes and incurs additional cost.

B Implementation Details

  • Classification uses the native activation with SiLU as the alternative, while segmentation uses the native activation with tanh as the alternative.These fixed task-level choices are used across datasets and architectures, following the activation-choice ablation.Source: B Implementation Details · p. 21
  • The main results use independent node-level masks with p = 0.5, giving both activations equal probability; channel- and layer-level masks are excluded.The passage states p = 0.5 maximizes Bernoulli-switch entropy and avoids systematic preference for either activation.Source: B Implementation Details · p. 21
  • Training matches each deterministic baseline’s optimizer, schedule, batch size, augmentation, and epochs, changing only activation switching.Classification uses cross-entropy and segmentation uses soft-Dice plus binary cross-entropy; Table 8 summarizes per-dataset hyperparameters and reports weight decay 5×10−4.Source: B Implementation Details · pp. 21–22, Table 8 · p. 22
  • Inference keeps trained weights fixed and uses T = 5 stochastic predictions, whose empirical mean and variance estimate the prediction and uncertainty.This increases per-input inference cost by approximately a factor of five; results use one random seed per method except the five-member Deep Ensemble.Source: B Implementation Details · pp. 21, 23
  • Experiments pair U-Net and Swin Transformer V2 with segmentation datasets, and ResNet-style or Swin Transformer V2 backbones with classification datasets.Each backbone matches its deterministic baseline, with only the activation block replaced by the stochastic switch.Source: B Implementation Details · p. 21
  • The evaluation uses official or standard dataset splits and designated cross-dataset OOD sets, including RoadTracer, CHASE-DB1, ISBI 2012, SVHN, and CIFAR-100.BloodMNIST and PathMNIST are treated as a cross-domain OOD pair.Source: B Implementation Details · p. 22
  • Training augmentation varies by task: segmentation uses crops, flips, and 90° rotations; classification uses crops and flips, with no test-time segmentation augmentation.High-resolution medical classification uses random resized crops and horizontal flips; inputs are channel-normalized using dataset-specific statistics.Source: B Implementation Details · p. 22
  • Baselines include MC Dropout, DDU, Masksembles, five-model Deep Ensembles, Laplace Redux, Test-Time Augmentation, and iterative segmentation refinement.MC Dropout uses rate 0.2 and T = 5 passes; the iterative segmentation baseline uses three refinement iterations.Source: B Implementation Details · pp. 22–23

C Mathematical Analysis of Stochastic Activation Switching

The appendix characterizes stochastic activation switching from unit-level randomness through full-network predictions and Monte Carlo estimators. It bounds per-unit variability and establishes convergence rates for empirical predictive summaries under stated conditions.

C.1 Notation and Setup

  • With trained weights fixed, stochastic masks generate predictions; repeated passes estimate their mean and activation-induced predictive variance.The analysis treats the current pre-activation as fixed when isolating a unit’s mask randomness, then averages over upstream randomness for full-network expectations.Source: §C.2 · pp. 23–24, §C.1 · p. 23

C.3 Per-Unit Moments and Bounded Activation Gap

  • For a fixed pre-activation, a unit’s output is a two-point random variable, and its variance is zero where the two activations agree.The per-unit variance is controlled by the switching probability and the activation gap; a bounded gap yields a bounded per-unit variance.Source: §C.3 · p. 24
  • Theorem C.1 bounds the noise envelope for one activation unit, but this bound does not directly bound full-network prediction variance.Later layers may amplify, damp, or combine per-unit perturbations.Source: §C.3 · p. 25

Define the gap function

  • For activation pairs with an unbounded global gap, the analysis gives a finite bound on any bounded pre-activation range.For ReLU–SiLU, the appendix gives an explicit global activation-gap constant used in the classification experiments.Source: §C.3 · p. 25

C.4 Convergence of the Empirical Estimators

  • The empirical mean is unbiased for the population mean, while the empirical variance estimator has finite-sample bias that vanishes as T increases.The appendix also notes that its concentration-bound constants are conservative, so inference-pass counts should balance uncertainty stability and computation.Source: §C.4 · pp. 27, 29
  • The empirical variance estimator converges to the population predictive variance, with root-mean-squared error O(T^-1/2).The analysis assumes bounded outputs and independently resampled masks across stochastic passes.Source: §C.4 · pp. 26, 28

D Computational Cost and Timing Analysis

Stochastic activation switching keeps training to one model while requiring repeated stochastic passes for uncertainty at inference. Its CIFAR-10 timing is close to MC Dropout and substantially below five-member ensemble training.

D Computational Cost and Timing Analysis

  • 42.73 ± 0.36s training and 14.21 ± 0.18s inference: our method is close to MC Dropout at T = 5, while five-member ensemble training takes 164.52s.A five-member ensemble requires five independently trained models; the method retains one trained network, reducing training and storage requirements.Source: D Computational Cost and Timing Analysis · p. 29, Table 9 · p. 30
  • Across classification and segmentation, the method retains a single-model training profile, while its multi-pass inference cost is generally comparable to MC Dropout and Deep Ensembles.Deterministic inference, DDU, Laplace Redux, and Masksembles generally have lower wall-clock inference latency.Source: D Computational Cost and Timing Analysis · pp. 29–30, Table 10 · p. 30
  • The main computational advantage is single-model training and storage, not minimum inference latency; repeated stochastic passes provide the predictive diversity used for uncertainty estimation.The timing protocol measures end-to-end procedures on one NVIDIA H100 GPU; our method and MC Dropout use five stochastic passes.Source: D Computational Cost and Timing Analysis · pp. 29–30

E Corruption Robustness

On CIFAR-10-C Gaussian noise, stochastic activation switching stays close to MC Dropout in accuracy while retaining intermediate information flow. Its uncertainty metrics are competitive across corruption severities.

E Corruption Robustness

  • Under mild corruption, our method preserves strong predictive accuracy and obtains the best Brier score, while MC Dropout has the highest accuracy at moderate and severe corruption.The text attributes MC Dropout’s heavier-corruption accuracy to feature masking as a noise-robust regularizer; that masking removes part of the intermediate representation.Source: E Corruption Robustness · p. 30, Figure 6 · p. 31
  • At moderate and severe corruption, stochastic activation switching remains close to MC Dropout in accuracy while keeping information flow mostly intact through activation switching.The evaluation uses CIFAR-10-C Gaussian noise at severities 1, 3, and 5.Source: E Corruption Robustness · p. 30
  • At the highest corruption severity, our method achieves the best NLL and ECE, while Laplace Redux has the best Brier score but substantially lower accuracy.The NLL and ECE markers are and, respectively.Source: E Corruption Robustness · p. 31

F OOD Uncertainty Distributions

Across the illustrated ID/OOD comparisons, OOD uncertainty is generally higher and more dispersed, though similar aerial-road imagery produces greater overlap. Segmentation examples show uncertainty localized to ambiguous structures and boundaries.

F OOD Uncertainty Distributions

  • ID samples usually have compact, low-uncertainty scores, while OOD scores shift higher and become more dispersed.The distributions overlap more for Massachusetts→RoadTracer, where both datasets contain aerial road imagery with similar visual structure.Source: F OOD Uncertainty Distributions · p. 31, Figure 7 · p. 32
  • For DRIVE→CHASE-DB1, CIFAR-10→SVHN, and CREMI→ISBI, OOD samples occupy a broader high-uncertainty range.For Massachusetts→RoadTracer, the OOD distribution still has a wider uncertainty spread despite stronger overlap.Source: F OOD Uncertainty Distributions · pp. 31–32
  • On ISBI 2012 with a model trained on CREMI, variance-based uncertainty maps concentrate around ambiguous membranes, boundaries, and structures where stochastic passes disagree.The figure presents each row as original image, ground-truth mask, prediction, and uncertainty map.Source: F OOD Uncertainty Distributions · p. 32, Figure 8 · p. 33

Related papers

Loading 2610.29875v1…