DESS: A Robust Uncertainty Layer for Embedding-Space Models

Abstract

Embedding-space systems usually retrieve, rank, or recommend from a single point representation, even when the task would benefit from uncertainty or controlled diversity. We introduce DESS, a lightweight uncertainty layer that augments an existing embedding model with a predicted mean vector and an independent per-dimension spread vector. DESS is trained with deterministic element-wise losses, requires no sampling during training, and maps inference-time samples back to available outputs by nearest-neighbor selection, leaving the trained predictor decoupled from a fixed output dictionary. The central question is whether such a simple diagonal uncertainty layer is meaningful and safe in high-dimensional embedding spaces. We show that the learned spread captures real local geometry: on 20 Newsgroups sentence-transformer embeddings, it is strongly rank-correlated with empirical within-class variation (Spearman $$\rho = 0.94$$ ). However, it is not epistemic uncertainty: the spread norm performs at chance for held-out-class OOD detection (AUROC 0.49). We then show that inference-time sampling is governed by concentration of measure, with sample displacement concentrating near the spread-vector norm. This yields a training-set geometric rule for choosing the spread multiplier $$\beta $$ , which predicts the safe sampling regime on AG News. On SQID/ESCI reranking, the strongest DESS variant preserves the frozen SBERT query embedding as the mean and learns only local spread; with $$\beta = 0.09$$ , sampled reranking matches the SBERT-text baseline without surpassing multimodal fusion or cross-encoder systems. DESS therefore provides a modular, geometry-aware uncertainty layer for embedding-space models, provided its spread is calibrated to local embedding geometry.