---
title: Shallow Feature Uncertainty Modeling (SUFM)
url: https://www.emergentmind.com/topics/shallow-feature-uncertainty-modeling-sufm
type: topic
---

# Shallow Feature Uncertainty Modeling (SUFM)

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 [2106.13066] [2507.17281].

## 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) [2504.12627] [2602.15747] [2202.03958] [2604.01305].

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" [2106.13066], SUFM is realized through shallow random features in a universal RKHS. Let $k$ be a positive definite kernel on $X \subset \mathbb{R}^p$ with RKHS $H$. Random features approximate the kernel by a finite-dimensional randomized feature map $\hat \phi : X \to \mathbb{R}^L$ such that
$$
k(x,x') \approx \frac{1}{L}\sum_{i=1}^L \hat\phi_i(x)\hat\phi_i(x').
$$
For stationary kernels, the paper uses Rahimi–Recht random Fourier features,
$$
\hat\phi_i(x)=\sqrt{\frac{2}{L}}\cos(a_i^\top x+b_i),\quad a_i\sim p(a),\quad b_i\sim \mathrm{Uniform}[0,2\pi],
$$
and for the Gaussian RBF kernel $k(x,x')=\exp(-\|x-x'\|_2^2/2\ell^2)$, one samples $a_i\sim \mathcal{N}(0,\ell^{-2}I_p)$. Random ReLU features,
$$
\hat\phi_i(x)=\max\{0,w_i^\top x+b_i\},
$$
are also reported as inducing universal RKHSs.

The dynamical system is modeled as
$$
x_{t+1}=h(x_t,u_t)+f(x_t,u_t),
$$
where $h$ is known and $f$ 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
$$
f(x,w)=\hat\phi(x)^\top w.
$$
With Gaussian prior $w\sim\mathcal{N}(0,I_L)$ and Gaussian observation noise, Bayesian linear regression yields
$$
\mu_{\omega\mid\mathcal D}=(\Phi^\top\Phi+\sigma^2 I_L)^{-1}\Phi^\top \hat y,\qquad
\Sigma_{\omega\mid\mathcal D}=\sigma^2(\Phi^\top\Phi+\sigma^2 I_L)^{-1},
$$
and predictive residual law
$$
\hat y(x)\mid\mathcal D\sim \mathcal N\!\left(\hat\phi(x)^\top \mu_{\omega\mid\mathcal D},\ \sigma^2+\hat\phi(x)^\top \Sigma_{\omega\mid\mathcal D}\hat\phi(x)\right).
$$
The resulting URF model is set-valued:
$$
x_{t+1}\in \mathcal F_{\mathrm{URF}}(x_t):=\{h(x_t)+\hat\phi(x_t)^\top w:\ w\in\mathcal W\},
$$
with ellipsoidal credible set
$$
\mathcal W=\left\{w\in\mathbb R^L:(w-\mu_{\omega\mid\mathcal D})^\top \Sigma_{\omega\mid\mathcal D}^{-1}(w-\mu_{\omega\mid\mathcal D})\le \chi^2_L(\alpha_\omega)\right\}.
$$

The distinctive contribution is not only predictive uncertainty but worst-case dynamics analysis. For horizon $N$, state recursion $x_{t+1}=h(x_t)+\hat\phi(x_t)^\top w_t$, and negative stage cost $\hat c(x_t)=-c(x_t)$, the paper studies minimization of cumulative negative cost over $w_t\in\mathcal W$. The Hamiltonian is affine in the shallow parameters,
$$
H(x,p,w)=\hat c(x)+p^\top \hat\phi(x)^\top w,
$$
so discrete-time PMP gives costate recursion and Hamiltonian minimization over $\mathcal W$. Because $\mathcal W$ is an ellipsoid, the minimizer is closed form, and the inexact update
$$
w_t^{k+1}=w_t^k+\gamma_k(s_t-w_t^k),\qquad
s_t=\arg\min_{w\in\mathcal W}\langle \nabla_{w_t}\hat J,w\rangle,
$$
is equivalent to a Frank–Wolfe step, with
$$
\nabla_{w_t}\hat J=\hat\phi(x_t)\,p_{t+1}.
$$
Choosing $\gamma_k=1$ recovers exact PMP. The rollout of the one-step map is thereby interpreted as a depth-$N$ computational graph whose adjoint equations are backpropagation through time.

The paper also exploits low-dimensional structure through random-feature nonlinear component analysis, projecting $\hat\phi(x)$ to $\hat\psi(x)=P\hat\phi(x)$ with $\hat L\ll L$. Computationally, BLR requires $O(TL^2)$ to compute $\Phi^\top\Phi$ and $O(L^3)$ to invert the $L\times L$ matrix, compared with Gaussian process regression’s $O(T^3)$. In experiments on a 2D source-spiral map, a Van der Pol oscillator, and a damped pendulum, training used $L=1000$ Fourier random features reduced to $\hat L=100$ 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 $\gamma_k=2/(k+2)$ yielded better optimization than full steps $\gamma_k=1$ 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" [2504.12627] describes DPOSE within SchNet as a direct instantiation of this idea: the final layer is replaced by $M$ parallel output heads, all earlier layers share weights, and each head produces an independent prediction $f_i(x)$. The ensemble mean and variance are computed analytically in one forward pass,
$$
\mu(x)=\frac{1}{M}\sum_{i=1}^M f_i(x),\qquad
\sigma^2(x)=\frac{1}{M}\sum_{i=1}^M f_i(x)^2-\mu(x)^2,
$$
and optimized with the heteroscedastic Gaussian negative log-likelihood
$$
L_{\mathrm{NLL}}=\frac{1}{2}\sum_n\left[\log \sigma_n^2+\frac{(y_n-\mu_n)^2}{\sigma_n^2}\right].
$$
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 $(0.09)$ to CCl4 $(62)$, from CH4 $(0.09)$ to SiH4 $(2.7)$, and from HF $(0.21)$ to HCl $(3.2)$; on OC20, intermetallic slabs show low uncertainty with upper whisker $<0.01$, while non-metals have median variance $\approx 0.03$; on the Gold dataset, elemental OOD from Au to Ag raises mean variance from $0.00022$ to $0.61304\ \mathrm{eV/atom}$.

"How to Train a Shallow Ensemble" [2602.15747] 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,
$$
\varepsilon_i=(w^{(k)})^\top h_i+b^{(k)},\qquad E^{(k)}(A)=\sum_i \varepsilon_i^{(k)},\qquad F_i^{(k)}(A)=-\nabla_{r_i}E^{(k)}(A).
$$
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 $SE_E$, yields miscalibrated force uncertainties; explicit joint energy-and-force NLL, denoted $SE_{E,F}$ with $L_{E,F}=L_E+L_F$, is reported as essential for reliable force calibration. Representative force RLL values for $SE_{E,F}$ are $59.6\%$ on BMIM, $43.8\%$ on H$_2$O, $69.3\%$ on BaTiO$_3$, $36.8\%$ on Ala4, and $5.4\%$ 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 $96\%$.

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" [2202.03958] introduces DSU, which treats per-channel feature mean and standard deviation as random variables. For a feature map $x\in\mathbb R^{B\times C\times H\times W}$,
$$
\mu_{b,c}=E_{h,w}[x_{b,c,h,w}],\qquad \sigma_{b,c}=\mathrm{Std}_{h,w}[x_{b,c,h,w}],
$$
and the statistic vector is $s_{b,c}=[\mu_{b,c},\sigma_{b,c}]^\top$. DSU assumes
$$
s_{b,c}\sim \mathcal N(m_{b,c},\Sigma_c),\qquad
\Sigma_c=\mathrm{diag}(\Sigma_{\mu,c}^2,\Sigma_{\sigma,c}^2),
$$
with batch-adaptive variances
$$
\Sigma_{\mu,c}^2=\frac{1}{B}\sum_b(\mu_{b,c}-E_b[\mu_{b,c}])^2,\qquad
\Sigma_{\sigma,c}^2=\frac{1}{B}\sum_b(\sigma_{b,c}-E_b[\sigma_{b,c}])^2.
$$
Sampling uses diagonal Gaussian noise and an AdaIN-style re-stylization,
$$
\mathrm{DSU}(x)=(\sigma+\epsilon_\sigma\Sigma_\sigma)\odot (x-\mu)/\sigma + (\mu+\epsilon_\mu\Sigma_\mu).
$$
No extra loss is added; the standard task loss is unchanged. Training applies DSU stochastically with probability $p$, reported as $0.5$ by default and $0.1$ 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 $79.0$; inserting DSU at positions $0$–$3$ yields $82.2$, positions $1$–$4$ yield $83.1$, positions $2$–$5$ yield $83.5$, and positions $0$–$5$ yield $84.1$. On GTA5$\rightarrow$Cityscapes semantic segmentation, the baseline mIoU is $37.0$, and the corresponding DSU variants reach $41.1$, $40.9$, $42.1$, and $43.1$. 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 $76.2$ to $73.4$ and further from $65.0$ to $64.1$ 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" [2604.01305] extends SUFM to sparse sensing and high-dimensional field reconstruction. SHRED maps a lagged sensor window $S=S(t-L+1:t)\in\mathbb R^{L\times p}$ to a spatial state $X$ through a recurrent encoder and shallow decoder; UQ-SHRED turns this into a conditional distribution learner by introducing Gaussian noise,
$$
X=G_\theta(S,z),\qquad z\sim \mathcal N(0,I_{d_z}),\qquad z\perp S,
$$
with the same $z$ replicated across the lag window for temporal coherence. The training objective is the multivariate energy score. The paper uses the practical two-sample loss
$$
L_{\mathrm{ES}}(\theta;S,X)=\frac{1}{2}\big(\|\hat x_1-X\|_2+\|\hat x_2-X\|_2\big)-\frac{1}{2}\|\hat x_1-\hat x_2\|_2,
$$
where $\hat x_j=G_\theta(S,z_j)$ and $z_1,z_2$ are i.i.d. Gaussian draws. Inference then resamples the same trained network $K$ 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, $50\%\rightarrow 53.9\%$, $70\%\rightarrow 71.5\%$, $90\%\rightarrow 87.3\%$, $95\%\rightarrow 90.8\%$, and $99\%\rightarrow 93.9\%$, with CRPS $\approx 0.343$; for solar activity, $95\%\rightarrow 94.6\%$ and CRPS $\approx 0.016$; for 1D rotating detonation ignition, calibration is shift-sensitive, with Run 0 giving $95\%\rightarrow 95.9\%$ and Run 1 giving $95\%\rightarrow 91.5\%$.

"Fully Automated SAM for Single-source Domain Generalization in Medical Image Segmentation" [2507.17281] 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 $f\in\mathbb R^{B\times C\times H\times W}$, 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
$$
f'=\gamma(x)\cdot (f-\mu(x))\cdot \sigma(x)^{-1}+\beta(x).
$$
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 $81.51$ and final prediction Dice $84.53$, whereas deeper insertions degrade performance, including position 5 with final Dice $77.25$. The combined disturbance mechanism (“United”) outperforms Gaussian only and Poisson only: preliminary/final Dice are $72.46/76.96$ for Gaussian, $78.93/82.53$ for Poisson, and $81.51/84.53$ for the combined variant. In the prostate benchmark, AGM only gives average Dice $81.31$, AGM+SUFM gives $83.62$, AGM+IPEF gives $83.52$, and AGM+SUFM+IPEF gives $84.53$; the full FA-SAM average is $84.53$ versus $79.54$ for SAMMed. On fundus vessel segmentation, FA-SAM reports average Dice $70.48$, above RAS$^4$DG $(69.37)$, DeSAM $(69.35)$, and SAMMed $(67.95)$.

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 [2106.13066] [2504.12627] [2202.03958] [2604.01305] [2507.17281]. 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 $L$ 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 [2106.13066]. 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 [2602.15747] [2504.12627]. 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 [2202.03958]. 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 [2604.01305]. 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 [2507.17281].

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 $u$, online BLR updates, and scalability through sparsity and feature selection [2106.13066]. 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 [2504.12627]. 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 [2602.15747]. 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 [2604.01305]. 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 [2202.03958].

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.

Source: https://www.emergentmind.com/topics/shallow-feature-uncertainty-modeling-sufm