Token-wise Variational Bayesian Inference (T-VBI)

• Token-wise Variational Bayesian Inference (T-VBI) recasts deterministic segmentation tokens as latent variables to facilitate uncertainty-aware predictions.
• It employs a two-phase approach by first estimating token statistics from unlabeled data and then using reparameterized Gaussian sampling at inference.
• By confining Bayesian reasoning to a low-dimensional token space, T-VBI preserves pretrained knowledge, reduces computational costs, and enhances interpretability in safety-critical settings.

Token-wise Variational Bayesian Inference (T-VBI) is a Bayesian adaptation mechanism introduced within E-BayesSAM for uncertainty-aware ultrasonic segmentation. Its central idea is to reinterpret the Segment Anything Model’s output tokens as dynamic probabilistic weights and to reparameterize them as latent variables without auxiliary training, thereby enabling training-free variational Bayesian inference in the token space rather than over SAM’s full parameter set (Huang et al., 24 Aug 2025). In this formulation, uncertainty estimation is attached directly to the decoder’s token outputs, with the stated aim of preserving pretrained knowledge, reducing computational burden, and supporting token-level interpretability in safety-critical medical applications.

1. Conceptual basis: output tokens as latent variables

T-VBI begins from a reinterpretation of SAM’s decoder outputs. In the E-BayesSAM formulation, the output tokens $\mathbf{T}$, ordinarily treated as deterministic vectors encoding segmentation-relevant features such as boundaries, are recast as dynamic weights that can take the place of fixed weights $W$ in conventional neural networks (Huang et al., 24 Aug 2025). The stated theoretical motivation is that, by the universal approximation theorem, the token generator $g(x;\theta)$ in SAM can approximate mappings from an input $x$ to $W$, so the output tokens can be viewed as input-dependent weights.

This reinterpretation changes the locus of Bayesian modeling. Rather than placing distributions over all of SAM’s millions of weights, T-VBI treats the token variables themselves as stochastic latent variables endowed with probability distributions. The immediate implication is architectural: Bayesian reasoning is confined to a low-dimensional token space, while the rest of the pretrained model remains unchanged. A common misunderstanding is to equate T-VBI with full-network Bayesian fine-tuning; in E-BayesSAM, the method is explicitly presented as an alternative to that strategy, motivated by the instability, cost, and retraining burden of Bayesian adaptation over the full parameter space (Huang et al., 24 Aug 2025).

2. Mathematical formulation

The T-VBI procedure has two phases. In the first phase, token statistics are estimated from unlabeled data. For each output token dimension, the empirical mean and standard deviation are computed as

$\mu_j = \frac{1}{N} \sum_{i=1}^N t_{ij}$

and

$$\sigma_j = \sqrt{\frac{1}{N} \sum_{i=1}^N (t_{ij} - \mu_j)^2},$$

where $t_{ij}$ is the $j$-th dimension of the token from the $i$-th data sample (Huang et al., 24 Aug 2025). In the second phase, these statistics are used at inference time. For a new sample $W$0, the mean is set by the current output tokens, $W$1, while the variance $W$2 is taken from the first phase.

Stochastic output tokens are then generated by reparameterized Gaussian sampling:

$W$3

At the distributional level, each output-token dimension is modeled as

$W$4

This Gaussian parameterization is the basis for Monte Carlo uncertainty estimation without modifying or retraining the bulk of the model. The formulation is therefore variational in the sense that latent stochastic structure is assigned to token dimensions, but operationally it is realized through token-statistic estimation and reparameterized sampling in inference rather than through a conventional end-to-end variational optimization of all model parameters.

3. Predictive inference and uncertainty estimation

Given sampled tokens, segmentation predictions are generated by projecting them against image features. For $W$5 stochastic samples, the $W$6-th mask prediction is

$W$7

where $W$8 is the projection of image features (Huang et al., 24 Aug 2025). The final prediction is