Skip to content
Open access

DESS: A Robust Uncertainty Layer for Embedding-Space Models

Aug 2026 · KI - Künstliche Intelligenz · 0 citations · 8 references

TL;DR

DESS is introduced, a lightweight uncertainty layer that augments an existing embedding model with a predicted mean vector and an independent per-dimension spread vector that provides a modular, geometry-aware uncertainty layer for embedding-space models, provided its spread is calibrated to local embedding geometry.

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.

Read PDF

Similar papers

Preprint Jul 2026

Covariance Last-Layer Ensembles: Function-Space Diversity for Efficient Uncertainty Quantification

A Last-Layer Ensemble (LLE), $K$ linear units on one shared frozen feature map, is an efficient single-pass approach to the disagreement-based epistemic uncertainty for out-of-distribution (OOD) detection. Its weakness is that members share the backbone gradient and can converge toward the same function, collapsing the inter-member diversity the signal depends on. Whether last-layer diversity can be restored, and what mitigates the collapse, is an open question. The weight-orthonormality defining Orthonormal Certificates (OC), the weight-orthonormal special case of the LLE, is only an indirect correction; it decorrelates the weights of the members, not their predictions. Here, we instead target the collapse directly in function space, with a Covariance Last-Layer Ensemble (cov-LLE) that places a direct covariance penalty on member activations. Cov-LLE restores the function-space diversity that weight-orthonormality cannot, and at matched $K$ recovers much of the diversity and calibration of a deep ensemble at $1\times$ backbone cost (in-distribution prediction variance $0.05\!\to\!9.3$ vs. $22.1$ ($\times10^{-3}$), and ECE $0.135\!\to\!0.090$ vs. $0.035$, for a $K\times$-cost deep ensemble), at no cost to accuracy. Viewing OC as a last-layer ensemble also organizes detectors into a two-axis taxonomy (by how their units are trained and how their outputs are scored) and exposes the OC score as a magnitude, motivating a scale-invariant, label-free direction score that repairs its near-OOD failure, adding $+0.16$ to $+0.18$ ROC AUC on every backbone.

H. M. Gillis, Isaac Xu, Gabriel Spadon et al. · 0 citations
Conference Open access 2026

Progressive Multi-Objective Optimization for Improved t-SNE Embeddings

A progressive multi-objective optimization framework is proposed that enhances t-SNE by integrating complementary loss functions, including a ranking-aware divergence (KLmax) and a Wasserstein-based term for global alignment.

S. Belhaouari, Skander Bensegueni, Lyes Fennour et al. · 0 citations
Book Open access Jul 2026

Scaling and Stabilizing Large-Scale Embedding-Based Retrieval

A unified pipeline deployed at Walmart that addresses both signal quality and model evolution is presented, and a Warm-Start Distillation technique that transfers domain-specific expertise from the legacy model to the new backbone is introduced.

Zhen Yang, Juexin Lin, Hongwei Shang et al. · 1 citation
Open access Aug 2026

Do Retrieval-Trained Embeddings Help Linear Contextual Bandits?

Text embeddings from retrieval-tuned (dual-encoder) models are increasingly used as context features in contextual bandits for recommendation, on the assumption that an embedding space optimized for inner-product similarity will speed up a linear exploration policy. This study tests that assumption with a controlled, shared-encoder design: the same BERT-base model in two forms, vanilla (mean-pooled) and retrieval-fine-tuned (MS MARCO dot-product), used as frozen bandit context. Experiments span three datasets (MIND, MovieLens-1M, Amazon CDs and Vinyl), two linear policies (LinUCB and linear Thompson sampling), and two dimensionality-reduction methods (PCA and random projection), over 20 seeds. Neither the training objective nor the reduction method determines performance on its own; cumulative regret is governed by their interaction, which is dataset-dependent. On MovieLens the reduction reverses the encoder ranking (PCA favors the retrieval-tuned encoder, random projection the vanilla one), with large effects in both directions, while on MIND the encoders are close. The reduction method can be the larger lever, moving up to 34% of the learnable margin. A ridge-regression probe on the candidate contexts screens the reduction choice offline, without running the bandit. For linear exploration policies, the encoder and the reduction should be treated as a joint choice rather than assuming retrieval-tuned embeddings are universally preferable.

Mustafa Canim · 0 citations
Preprint Aug 2026

Beyond a Bag of Features: Set-Level Instability in Sparse Autoencoders

It is found that SAE activation sets do not recover human category boundaries or within-category typicality more faithfully than dense embeddings or residual-stream states, but instead track model-internal similarity structure.

Nikolai Bolik, Lennart Stöpler, Artur Andrzejak · 0 citations
Preprint Aug 2026

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Optimizer-state quantization is commonly designed for Adam's dense, parameter-aligned first- and second-moment arrays. This abstraction breaks for memory-efficient optimizers, whose states may be factored, confidence-modulated, or maintained in a projected space, so similar reconstruction error can produce different update error. We formulate optimizer-state quantization as a joint problem over representation, topology, and update semantics. We then introduce Adaptive Log-Space (AL) quantization for non-negative states. AL fits each block's observed nonzero logarithmic interval and reserves a separate code for exact zero, enforcing $q = 0 \Leftrightarrow x = 0$; signed momentum and state precision remain independently selectable. Controlled probes show that adaptive ranges reduce update error and temporal drift, exact-zero reservation preserves dormant states, and state topology constrains useful block granularity. End-to-end language-model training evaluates the resulting policy across dense, factored, confidence, and projected optimizer states. On TinyLlama-1.1B, AL8 with uniform 8-bit momentum reaches 72.90 perplexity versus 73.54 for bitsandbytes 8-bit AdamW, with comparable optimizer-state storage and higher throughput. CAME matches reference-level final perplexity across three seeds when its non-negative states use AL16, while a semantic grouping-and-protection policy closes most of quantized Adafactor's 100K-step late-loss gap. These results make state topology and update semantics first-class design constraints for optimizer quantization.

Yan Wang · 0 citations