Epistemic Uncertainty Quantification for Pre-trained VLMs via Riemannian Flow Matching

Technical Summary: REPVLM – Epistemic Uncertainty Quantification for Pre-trained VLMs via Riemannian Flow Matching

1. Problem Statement

Pre-trained Vision-Language Models (VLMs) such as CLIP, BLIP, and SigLIP have achieved remarkable success in cross-modal representation learning through large-scale contrastive pre-training. However, these models are inherently deterministic, mapping inputs to single, fixed points in a high-dimensional embedding space. This architecture lacks an intrinsic mechanism to quantify epistemic uncertainty (the model's ignorance or lack of knowledge regarding its own representations).

While existing uncertainty quantification (UQ) frameworks for VLMs exist, they predominantly address aleatoric uncertainty (data ambiguity) by learning probabilistic embeddings. Methods designed for epistemic uncertainty, such as deep ensembles or Monte Carlo Dropout, are often computationally prohibitive for large-scale pre-trained models or yield suboptimal, batch-dependent estimates. There is a critical need for a scalable, intrinsic metric that directly quantifies a model's confidence in its representations without requiring expensive retraining or parameter-space sampling.

2. Methodology: REPVLM

The authors propose REPVLM, a framework that quantifies epistemic uncertainty by estimating the probability density of VLM embeddings directly on their native manifold.

Theoretical Motivation
The paper establishes a theoretical link between embedding density and epistemic uncertainty through a three-step chain:

1. Covariance to Jacobian Norm: Under standard Bayesian assumptions and local linearity, high epistemic uncertainty corresponds to a large Parameter-Jacobian norm, indicating sensitivity to parameter perturbations.
2. Jacobian Norm to Data Density: In high-density regions of the training data, the encoder's Jacobian norm tends to be small (smoothness), whereas in low-density (Out-of-Distribution) regions, it remains unconstrained.
3. Data Density to Embedding Density: Motivated by alignment-and-uniformity dynamics in contrastive learning, the authors posit an empirical assumption that the density of the input data $p(x)$ correlates with the density of the resulting embeddings $p(z)$ on the hypersphere.
   Consequently, the negative log-likelihood, $-\log p(z)$, serves as a principled proxy for epistemic uncertainty: low-density regions signify model ignorance.

Riemannian Flow Matching (RFM)
To compute $p(z)$ efficiently on the hypersphere $S^{d-1}$ (where $\ell_2$-normalized VLM embeddings reside), REPVLM utilizes Conditional Riemannian Flow Matching (CRFM):

• Manifold-Native Modeling: Unlike standard Euclidean Flow Matching, REPVLM operates directly on the hypersphere. It learns a vector field $v_t$ that transports a simple uniform base distribution $P_0 = \text{Unif}(S^{d-1})$ to the empirical modality-specific distributions $P_1$ (Image and Text).
• Geodesic Interpolation: The probability path is constructed along geodesics (great circles) rather than straight lines, respecting the intrinsic geometry of the semantic space.
• Unified Conditional Framework: A single neural network learns the vector field conditioned on the modality ($c \in \{\text{image}, \text{text}\}$), allowing for task-agnostic uncertainty estimation.
• Likelihood Computation: The exact log-density is computed by integrating the divergence of the learned vector field along the reverse ODE trajectory. To handle high-dimensional embeddings efficiently, the authors employ a manifold-constrained Hutchinson trace estimator to approximate the divergence with minimal computational overhead.

3. Key Contributions

• Theoretical Motivation: The paper provides a formal derivation demonstrating that $-\log p(z)$ is a monotone proxy for epistemic uncertainty in contrastive encoders, linking Bayesian uncertainty to embedding density.
• Manifold-Native Method (REPVLM): The authors introduce a scalable UQ method that extends Riemannian Flow Matching to VLM embedding spaces, enabling intrinsic probability density computation for both visual and textual modalities without modifying the pre-trained backbone.
• Empirical Validation: The method is validated across multiple benchmarks, demonstrating near-perfect correlation between uncertainty scores and prediction errors, significantly outperforming existing baselines in selective classification and OOD detection.

4. Experimental Results

The authors evaluated REPVLM on six downstream benchmarks (ImageNet-1K, Food101, CIFAR-100, ObjectNet, ImageNet-R, ImageNet-Sketch) using three different proxy datasets for training (Conceptual Captions, DataComp-1B, LAION-2B).

• Selective Classification: REPVLM achieved near-perfect Spearman correlation ($S \ge 0.988$) between uncertainty and prediction error across standard benchmarks. It consistently outperformed baselines like ProbVLM, Monte Carlo Dropout, and Embedding Norm in identifying model errors.
• Robustness to Distribution Shift: On distribution-shifted datasets (ObjectNet, ImageNet-R, ImageNet-Sketch), REPVLM was the only method to achieve $S=1.000$ across all shift cells, whereas other methods (e.g., ProbVLM, GMM) showed high variance or negative correlations.
• Applications:
  • OOD Detection: The uncertainty scores effectively separated In-Distribution (ID) data from Near-OOD and Far-OOD data, yielding high ROC and Precision-Recall curves.
  • Automated Data Curation: High-uncertainty samples flagged by REPVLM corresponded to noisy, blurry, or nonsensical inputs, demonstrating utility in cleaning large-scale web datasets.
• Efficiency: REPVLM requires only ~0.55 GFLOPs of additional overhead beyond the backbone, making it significantly more efficient than Monte Carlo Dropout (which requires ~10x forward passes) while providing more reliable epistemic estimates than ProbVLM.

5. Significance and Limitations

Significance
The paper argues that on-manifold embedding density is the operative quantity for epistemic uncertainty in pre-trained VLMs. By leveraging the geometric structure of the embedding space via Riemannian Flow Matching, REPVLM provides a scalable, intrinsic, and highly reliable metric for model confidence. This enables practical applications such as automated data curation and robust OOD detection without the computational cost of ensemble methods.

Limitations
The authors explicitly acknowledge several limitations:

1. Theoretical Proxy: The link between Bayesian epistemic uncertainty and embedding density is a motivating proxy argument based on empirical assumptions (specifically regarding density preservation in non-invertible contrastive encoders), not a formal identity.
2. Proxy Dependency: Like all post-hoc methods, REPVLM requires a proxy dataset representative of the target VLM's pre-training domain. Specialized VLMs may require in-domain proxies.
3. Geometric Assumptions: The theory assumes $\ell_2$-normalized embeddings on $S^{d-1}$ from cosine-similarity-trained encoders (like CLIP). The authors note that results may not extrapolate to models with different geometries (e.g., SigLIP, which showed degraded performance on in-distribution benchmarks due to texture-driven density variations).
4. Fairness Risks: The authors caution that density-based uncertainty conflates rarity with ignorance. Rare but valid inputs (e.g., from underrepresented groups) may be flagged as high-uncertainty outliers. In safety-critical settings, such samples should be candidates for human review rather than automatic removal.