Recovering Event Probabilities from Large Language Model Embeddings via Axiomatic Constraints

TL;DR

This work addresses the problem that LLM-derived event probabilities are often incoherent with the axioms of probability. It proposes an unsupervised, two-step VAE approach that learns a low-dimensional latent space from LLM embeddings and enforces the additivity constraint for complementary events by modifying a single latent variable (interpreted as log-odds) via a sign flip, yielding recovered probabilities $P_ ext{recovered}=rac{e^{ extbf{z}^{(1)}}}{1+e^{ extbf{z}^{(1)}}}$. Training interleaves standard reconstruction with a second objective to predict embeddings of complementary events, guided by a centered Gaussian prior ($p( extbf{z})= ext{N}( extbf{0}, extbf{I})$) and a $eta$-VAE regime ($eta=5$). Empirical evaluation on dice-related events shows recovered probabilities are more coherent than text-based judgments and closely track true probabilities, with ablation demonstrating the necessity of Step 2; results also indicate model- and layer-dependent performance, with Gemma-2-9b-instruct outperforming Llama-3.1-8b-instruct. The findings reveal that coherent probabilistic structure can be recovered from embeddings and suggest avenues for integrating axiomatic constraints into LLMs to enhance probabilistic reasoning in uncertain domains.

Abstract

Rational decision-making under uncertainty requires coherent degrees of belief in events. However, event probabilities generated by Large Language Models (LLMs) have been shown to exhibit incoherence, violating the axioms of probability theory. This raises the question of whether coherent event probabilities can be recovered from the embeddings used by the models. If so, those derived probabilities could be used as more accurate estimates in events involving uncertainty. To explore this question, we propose enforcing axiomatic constraints, such as the additive rule of probability theory, in the latent space learned by an extended variational autoencoder (VAE) applied to LLM embeddings. This approach enables event probabilities to naturally emerge in the latent space as the VAE learns to both reconstruct the original embeddings and predict the embeddings of semantically related events. We evaluate our method on complementary events (i.e., event A and its complement, event not-A), where the true probabilities of the two events must sum to 1. Experiment results on open-weight language models demonstrate that probabilities recovered from embeddings exhibit greater coherence than those directly reported by the corresponding models and align closely with the true probabilities.

Figures (6)

• Figure 1: A schematic illustration of the two-step VAE-based learning algorithm. The objective is to distill axiomatic constraints between embeddings into the modified latent variables. In the first step, a standard VAE is used to disentangle explanatory factors within the latent space from the model embeddings. This involves learning to reconstruct the original LLM embeddings, $\mathbf{e}$, by encoding them into a compressed latent space, $\mathbf{z}$, and subsequently decoding reconstructed embeddings, $\Tilde{\mathbf{e}}$. The latent space is modeled as a multivariate Gaussian distribution with mean $\mu$ and standard deviation $\sigma$. In the second step, a subset of the latent variables is modified after encoding, and the modified latent vector is passed through the decoder to predict the embeddings of complementary events, $\neg \Tilde{\mathbf{e}}$. The probabilistic encoder and probabilistic decoder are denoted $q_\phi(\mathbf{z} |\mathbf{e})$ and $p_\theta(\mathbf{e} |\mathbf{z})$ respectively. Owing to the symmetry property (i.e., $\neg (\neg \mathbf{e}) = \mathbf{e}$), both $\mathbf{e}$ and $\neg \mathbf{e}$ can be used as input embeddings.
• Figure 2: Comparison of event probability estimates elicited from Gemma-2-9b-instruct: (a) judged probabilities, (b) recovered probabilities, (c) recovered probabilities with Step 2 ablated, and (d) probabilities predicted by a linear probe. As indicated by the dotted black lines in the left panel, coherent probability estimates for event A and its complement (not-A) should sum to 1. The right panel compares the elicited probabilities with the true probabilities. Error bars represent standard errors of the mean for the window-binned scatter points.
• Figure 3: Visualization of the mean values of the 10 latent variables. Each panel compares the mean of the latent variables for embeddings corresponding to event A (horizontal axis) with those for embeddings corresponding to event not-A (vertical axis). Only the modified latent variable 1 exhibits a negative relationship.
• Figure 4: Relationship between judged probabilities (horizontal axis) and recovered probabilities (vertical axis). The dashed black lines indicate perfect agreement.
• Figure 5: Step 2 ablated. Visualization of the mean values of the 10 latent variables. Each panel compares the mean of the latent variables for embeddings corresponding to event A (horizontal axis) with those for embeddings corresponding to event not-A (vertical axis).