Task-Awareness Improves LLM Generations and Uncertainty

TL;DR

This paper introduces a task-aware, Bayesian decision-theoretic framework for LLMs that models outputs in a latent structure, improving response quality and uncertainty estimation across various tasks.

Contribution

It proposes a novel approach that incorporates latent structural modeling and Bayesian optimization to enhance LLM decoding and uncertainty quantification.

Findings

Bayes-optimal responses outperform standard decoding methods.

Uncertainty quantification aligns better with output correctness.

Framework is applicable across multiple structured tasks.

Abstract

In many applications of LLMs, natural language responses often have an underlying structure such as representing discrete labels, numerical values, or graphs. Yet, existing decoding and uncertainty estimation methods operate only in language space and largely disregard structural information. We address this by modeling LLM outputs directly in a task-dependent latent structure. By equipping this structure with a dissimilarity measure, we can compute Bayes-optimal responses. These are not selected from sampled generations but are newly synthesized by combining individual responses in the latent space. Across different tasks, Bayes-optimal responses consistently outperform standard decoding methods like beam search. Moreover, quantifying uncertainty via the induced Bayesian risk captures variations in terms of the latent structure and improves alignment with output quality and correctness. Our decision-theoretic framework is applicable to any problem that admits a latent response structure and enables reliable task-aware LLM predictions.