Task-Awareness Improves LLM Generations and Uncertainty

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.

In many downstream tasks that use LLMs, the generated free-form text has an underlying structure. For example, when asking an LLM to rate a text snippet for a certain criterion, the generated response may correspond to a number or a semantic class. Current decoding and uncertainty quantification methods discard this structure and only treat responses as simple strings. We look beyond the exact wording of responses and modeled LLM answers in a task-aware manner. This means we extract the task-relevant information of the free-form answer and disregard the rest, keeping only the part that really matters for the application. This allows us to (1) combine different responses a model gave for the same question into one final response that takes into account the model's belief and (2) better quantify how certain the model is in this response. We regard the model as uncertain not just when its beliefs over different response strings are uncertain, but only when their corresponding meanings differ strongly. This shift in perspective greatly improves the quality of the LLM's answers and makes estimates of the LLM’s uncertainty more accurate. We think this idea of interpreting free-form LLM responses in terms of their task-dependent meaning is a powerful paradigm that can be applied to many other use cases as well.