Shallow Feature Uncertainty Modeling (SUFM)
- SUFM is a strategy that concentrates uncertainty in shallow predictor components (e.g., random-feature weights or readout heads) rather than throughout the full network.
- It employs various formulations such as Bayesian random-feature dynamics, shallow ensembles, and DSU to improve calibration, robustness, and out-of-distribution generalization.
- Applications span dynamical systems, graph neural networks, and medical segmentation, demonstrating actionable gains in computational efficiency and reliable uncertainty estimates.
Shallow Feature Uncertainty Modeling (SUFM) denotes a family of uncertainty-aware modeling strategies in which uncertainty is concentrated in shallow components of a predictor—such as random-feature weights, final readout heads, shallow feature statistics, or a shallow decoder—rather than distributed across the full parameterization. In the works associated with this label, SUFM appears both as explicit terminology and as an inferred unifying principle: uncertainty-aware random features for dynamical systems place a posterior and ambiguity set on shallow parameters, shallow ensembles for graph and atomistic models diversify only the last layer, DSU-style domain generalization randomizes shallow feature statistics, UQ-SHRED injects stochasticity into a single shallow recurrent-decoder architecture, and FA-SAM uses a dedicated SUFM module in a shallow prompt-generation branch for automated medical segmentation (Agudelo-España et al., 2021, Zhuo et al., 23 Jul 2025).
1. Conceptual scope and recurring design pattern
The literature does not use a single canonical implementation of SUFM. Instead, the same structural idea recurs under different names: the Uncertainty-aware Random Feature (URF) dynamics model in universal RKHSs, Direct Propagation of Shallow Ensembles (DPOSE) in SchNet, shallow ensembles and last-layer Laplace approximations for interatomic potentials, DSU as uncertainty over shallow feature statistics, UQ-SHRED as distributional learning for sparse sensing, and the SUFM module inside FA-SAM’s Auto-prompted Generation Model (AGM) (Vinchurkar et al., 17 Apr 2025, Schäfer et al., 17 Feb 2026, Li et al., 2022, Gao et al., 1 Apr 2026).
Across these formulations, the uncertainty-bearing object is shallow, while the rest of the system is shared, frozen, lightweight, or interpreted as a downstream composition of one-step shallow maps. This commonality is architectural rather than tied to a single probabilistic formalism.
| Formulation | Shallow uncertainty carrier | Shared or fixed component |
|---|---|---|
| URF dynamics | BLR posterior and ellipsoidal credible set over random-feature weights | Nominal dynamics and rollout graph |
| DPOSE / shallow ensembles | Multiple last-layer heads or committee members | SchNet or atomistic backbone |
| DSU / FA-SAM | Sampled shallow feature statistics | Remaining encoder and segmentation pipeline |
| UQ-SHRED | Input noise through a single shallow recurrent decoder | One shared SHRED architecture |
A useful synthesis is that SUFM treats uncertainty as a structured perturbation of a low-dimensional or late-stage representation rather than as a full-network posterior. In some papers this supports tractable robust optimization; in others it yields single-pass uncertainty estimates, calibration improvements, or better domain generalization.
2. Random-feature dynamics, RKHS approximation, and worst-case analysis
In "Shallow Representation is Deep: Learning Uncertainty-aware and Worst-case Random Feature Dynamics" (Agudelo-España et al., 2021), SUFM is realized through shallow random features in a universal RKHS. Let be a positive definite kernel on with RKHS . Random features approximate the kernel by a finite-dimensional randomized feature map such that
For stationary kernels, the paper uses Rahimi–Recht random Fourier features,
and for the Gaussian RBF kernel , one samples . Random ReLU features,
are also reported as inducing universal RKHSs.
The dynamical system is modeled as
where 0 is known and 1 is an unknown smooth function in a universal RKHS. The paper focuses on the autonomous single-output case for exposition and approximates the unknown residual by the shallow linear-in-parameter model
2
With Gaussian prior 3 and Gaussian observation noise, Bayesian linear regression yields
4
and predictive residual law
5
The resulting URF model is set-valued:
6
with ellipsoidal credible set
7
The distinctive contribution is not only predictive uncertainty but worst-case dynamics analysis. For horizon 8, state recursion 9, and negative stage cost 0, the paper studies minimization of cumulative negative cost over 1. The Hamiltonian is affine in the shallow parameters,
2
so discrete-time PMP gives costate recursion and Hamiltonian minimization over 3. Because 4 is an ellipsoid, the minimizer is closed form, and the inexact update
5
is equivalent to a Frank–Wolfe step, with
6
Choosing 7 recovers exact PMP. The rollout of the one-step map is thereby interpreted as a depth-8 computational graph whose adjoint equations are backpropagation through time.
The paper also exploits low-dimensional structure through random-feature nonlinear component analysis, projecting 9 to 0 with 1. Computationally, BLR requires 2 to compute 3 and 4 to invert the 5 matrix, compared with Gaussian process regression’s 6. In experiments on a 2D source-spiral map, a Van der Pol oscillator, and a damped pendulum, training used 7 Fourier random features reduced to 8 by random-feature PCA; the uncertainty tube widened under scarce data and shrank under large data, while worst/best cost envelopes contained the true cost and moved toward the mean as training size increased. The paper further reports that the Frank–Wolfe schedule 9 yielded better optimization than full steps 0 or fixed small steps.
3. Shallow ensembles and probabilistic readouts in graph and atomistic models
In graph neural networks and machine-learned interatomic potentials, SUFM is implemented by making only the final readout stochastic or ensemble-valued while sharing the deep representation. "Uncertainty Quantification in Graph Neural Networks with Shallow Ensembles" (Vinchurkar et al., 17 Apr 2025) describes DPOSE within SchNet as a direct instantiation of this idea: the final layer is replaced by 1 parallel output heads, all earlier layers share weights, and each head produces an independent prediction 2. The ensemble mean and variance are computed analytically in one forward pass,
3
and optimized with the heteroscedastic Gaussian negative log-likelihood
4
For QM9, the model is trained end-to-end with 64 heads; for OC20 fine-tuning, all weights except the modified final layer are frozen; for Gold Molecular Dynamics, a SchNet checkpoint pretrained on OC20 is fine-tuned using NLL. The reported uncertainty behavior is explicitly out-of-domain sensitive: on QM9, example variances jump from CF4 5 to CCl4 6, from CH4 7 to SiH4 8, and from HF 9 to HCl 0; on OC20, intermetallic slabs show low uncertainty with upper whisker 1, while non-metals have median variance 2; on the Gold dataset, elemental OOD from Au to Ag raises mean variance from 3 to 4.
"How to Train a Shallow Ensemble" (Schäfer et al., 17 Feb 2026) systematizes this regime for atomistic machine learning. Here all ensemble members share the full feature backbone and differ only in the final linear readout layer,
5
Committee means and variances are defined across members for total energy and for each force component. The paper emphasizes that explicit NLL optimization improves calibration relative to random last-layer sampling or a last-layer Laplace approximation, and that force uncertainty must be modeled explicitly. Energy-only training, denoted 6, yields miscalibrated force uncertainties; explicit joint energy-and-force NLL, denoted 7 with 8, is reported as essential for reliable force calibration. Representative force RLL values for 9 are 0 on BMIM, 1 on H2O, 3 on BaTiO4, 5 on Ala4, and 6 on Carbon. The paper further reports that a shallow ensemble of approximately 32 members is sufficient, that full-model fine-tuning from an energy-probabilistic initialization or an LLPR sample yields negligible reduction in calibration quality relative to training from scratch, and that training time can be reduced by up to 7.
Taken together, these papers show that SUFM in graph and atomistic models is not merely a cheap approximation to a deep ensemble. It is a design in which epistemic uncertainty is deliberately localized in the readout space, sometimes with frozen deep features, sometimes with end-to-end NLL so that uncertainty information can flow back into the backbone. A plausible implication is that this localization works best when representation learning is already strong and the primary uncertainty lies in mapping shared features to target observables.
4. Uncertain shallow feature statistics for out-of-distribution generalization
A second major SUFM lineage places uncertainty not on readout weights but on shallow feature statistics themselves. "Uncertainty Modeling for Out-of-Distribution Generalization" (Li et al., 2022) introduces DSU, which treats per-channel feature mean and standard deviation as random variables. For a feature map 8,
9
and the statistic vector is 0. DSU assumes
1
with batch-adaptive variances
2
Sampling uses diagonal Gaussian noise and an AdaIN-style re-stylization,
3
No extra loss is added; the standard task loss is unchanged. Training applies DSU stochastically with probability 4, reported as 5 by default and 6 for large-scale ImageNet training, while inference disables DSU.
The paper’s ablations are directly relevant to SUFM because they isolate layer depth. On PACS classification, the baseline average accuracy is 7; inserting DSU at positions 8–9 yields 0, positions 1–2 yield 3, positions 4–5 yield 6, and positions 7–8 yield 9. On GTA50Cityscapes semantic segmentation, the baseline mIoU is 1, and the corresponding DSU variants reach 2, 3, 4, and 5. The interpretation offered in the paper is that shallow statistics encode style or domain characteristics such as color, contrast, and texture, so randomizing them exposes the network to plausible domain shifts during training. The same mechanism improves Office-Home classification, person re-identification, and ImageNet-C robustness, with corruption mCE reduced from 6 to 7 and further from 8 to 9 when combined with APR.
This line of work establishes a version of SUFM in which the uncertainty object is a shallow normalization-like statistic rather than a weight posterior or ensemble disagreement. The method is parameter-free, relies on mini-batch estimates of uncertainty, and is explicitly motivated by out-of-distribution generalization rather than calibrated posterior inference.
5. Distributional shallow decoders and automated prompting
"UQ-SHRED: uncertainty quantification of shallow recurrent decoder networks for sparse sensing via engression" (Gao et al., 1 Apr 2026) extends SUFM to sparse sensing and high-dimensional field reconstruction. SHRED maps a lagged sensor window 00 to a spatial state 01 through a recurrent encoder and shallow decoder; UQ-SHRED turns this into a conditional distribution learner by introducing Gaussian noise,
02
with the same 03 replicated across the lag window for temporal coherence. The training objective is the multivariate energy score. The paper uses the practical two-sample loss
04
where 05 and 06 are i.i.d. Gaussian draws. Inference then resamples the same trained network 07 times and forms means, variances, and coordinate-wise quantiles; the paper states a theorem on Monte Carlo quantile consistency. Reported observed coverages are close to nominal in several settings: for sea-surface temperature, 08, 09, 10, 11, and 12, with CRPS 13; for solar activity, 14 and CRPS 15; for 1D rotating detonation ignition, calibration is shift-sensitive, with Run 0 giving 16 and Run 1 giving 17.
"Fully Automated SAM for Single-source Domain Generalization in Medical Image Segmentation" (Zhuo et al., 23 Jul 2025) uses the term SUFM explicitly inside FA-SAM. The SUFM module is inserted into the shallow encoder layers of AGM, whose backbone is Meganet, to improve the quality of automatically generated box prompts under domain shift. For shallow feature maps 18, the module computes per-instance, per-channel means and variances, estimates batch-level uncertainty of these statistics, samples Gaussian perturbations, injects Poisson noise to model discrete imaging noise, and re-normalizes features as
19
The resulting AGM prediction is converted into a binary mask, the largest connected component is selected by BFS, and its bounding box is passed to SAM as a prompt. SAM’s image encoder and prompt encoder are frozen, while an Image-Prompt Embedding Fusion (IPEF) module is inserted into the mask decoder to combine multiscale image embeddings and prompt embeddings.
The paper reports that inserting SUFM between the first two convolutional blocks (“Layer 0–1”) is best, with preliminary prediction Dice 20 and final prediction Dice 21, whereas deeper insertions degrade performance, including position 5 with final Dice 22. The combined disturbance mechanism (“United”) outperforms Gaussian only and Poisson only: preliminary/final Dice are 23 for Gaussian, 24 for Poisson, and 25 for the combined variant. In the prostate benchmark, AGM only gives average Dice 26, AGM+SUFM gives 27, AGM+IPEF gives 28, and AGM+SUFM+IPEF gives 29; the full FA-SAM average is 30 versus 31 for SAMMed. On fundus vessel segmentation, FA-SAM reports average Dice 32, above RAS33DG 34, DeSAM 35, and SAMMed 36.
These two papers broaden SUFM beyond last-layer epistemic uncertainty. In UQ-SHRED, shallow uncertainty is a conditional distribution over high-dimensional reconstructions induced by stochastic input features. In FA-SAM, shallow uncertainty is a domain-shift simulator acting on prompt-generation features. In both cases the uncertainty mechanism is localized, but the output is operationally important: calibrated spatial confidence bands in one case and robust automated prompts in the other.
6. Limitations, interpretive cautions, and open directions
Across the cited works, “shallow” usually refers to the uncertainty-bearing component rather than to the full predictor. In the RKHS dynamics paper, the one-step model is shallow but the rollout is a multi-layer network under composition; in DPOSE and shallow ensembles, the backbone may be deep and only the last layer varies; in DSU and FA-SAM, the uncertainty mechanism acts in early layers of otherwise deep encoders; in UQ-SHRED, a shallow decoder sits behind a recurrent encoder (Agudelo-España et al., 2021, Vinchurkar et al., 17 Apr 2025, Li et al., 2022, Gao et al., 1 Apr 2026, Zhuo et al., 23 Jul 2025). This matters because SUFM is best understood as a localization principle for uncertainty rather than as a blanket preference for shallow architectures.
The main limitations are also localized. The RKHS/random-feature formulation assumes smooth dynamics and costs, linear-in-parameter Hamiltonian structure, Gaussian priors and observation noise, and an ellipsoidal credible region whose calibration may be conservative or may miss heavy tails; finite 37 and random-feature PCA introduce approximation error, and the paper notes that shallow random features may struggle with composite or deep structures because of depth-separation phenomena (Agudelo-España et al., 2021). In shallow ensembles for interatomic potentials, diversity alone is insufficient; energy-only objectives miscalibrate force uncertainties, last-layer Laplace approximations can exhibit Hessian spectral mismatch, and structurally similar OOD cases can remain hard to separate, as in amorphous-versus-bulk gold at low energies (Schäfer et al., 17 Feb 2026, Vinchurkar et al., 17 Apr 2025). In DSU, the Gaussian and diagonal assumptions ignore cross-channel covariance, estimation depends on mini-batch variance, and very small batch sizes can make the uncertainty estimate degenerate (Li et al., 2022). In UQ-SHRED, extreme sensor sparsity or large initial-condition shifts can lead to undercoverage, and the Gaussian latent noise may be misspecified for multimodal conditionals (Gao et al., 1 Apr 2026). In FA-SAM, poor box prompts are reduced but not eliminated, and performance is highly sensitive to placing SUFM early rather than deep in the encoder (Zhuo et al., 23 Jul 2025).
The open directions reported in these papers are correspondingly specific. For random-feature SUFM, proposed extensions include adaptive feature sampling, richer priors, non-Gaussian likelihoods, variational Bayesian treatments, integration with robust or stochastic control and MPC, extension to controlled systems with explicit 38, online BLR updates, and scalability through sparsity and feature selection (Agudelo-España et al., 2021). For shallow ensembles in graph and atomistic models, the cited directions include active learning, explicit aleatoric modeling, partial unfreezing of backbones, latent-distance signals, and domain-adversarial regularization (Vinchurkar et al., 17 Apr 2025). For atomistic calibration, the practical recommendation is to initialize from an energy-probabilistic shallow ensemble or sampled LLPR committee and then full-model fine-tune with joint energy-and-force NLL (Schäfer et al., 17 Feb 2026). For UQ-SHRED, suggested extensions include conformal prediction, richer conditionals such as mixture-of-experts, conditional normalizing flows, or diffusion/score-based decoders, and physics-aware constraints in the decoder (Gao et al., 1 Apr 2026). For DSU-style shallow statistic perturbation, the paper notes potential extensions such as per-layer adaptive application probabilities, learnable scales on batch-derived uncertainty, and mixture-of-Gaussians uncertainty for multimodal style sources (Li et al., 2022).
A plausible synthesis is that SUFM is most effective when uncertainty is dominated by readout ambiguity, shallow appearance statistics, or an explicit information bottleneck between sparse observations and high-dimensional targets. When the principal error source lies deeper in the representation itself, the literature repeatedly moves from purely shallow uncertainty toward end-to-end NLL training, partial or full fine-tuning, or hybridization with deeper adaptation mechanisms.