Technical Summary: Task-Awareness Improves LLM Generations and Uncertainty

Problem Statement

Large Language Models (LLMs) are increasingly deployed in applications where the output is not merely free-form text but possesses an underlying task-specific structure, such as discrete labels, numerical values, sets of entities, or graphs. Existing decoding strategies (e.g., beam search, greedy decoding) and uncertainty estimation methods operate primarily in the raw language space ($V^\infty$). These methods largely disregard the structural information inherent to the task. Consequently, they often fail to synthesize optimal responses that combine information across the model's full output distribution and provide uncertainty estimates that are misaligned with task-specific notions of correctness and quality.

Methodology

The authors propose a decision-theoretic framework that shifts LLM generation and uncertainty quantification from the raw language space to a task-dependent latent structure ($\mathcal{L}$).

1. Latent Space Mapping

The framework defines an encoding function $g_T: V^\infty \to \mathcal{L}$ that maps free-form LLM responses $s$ to a latent representation $\ell$. This latent space is tailored to the specific task:

• Classification: A discrete set of classes $[K]$.
• Multi-Answer QA: A power set of classes (represented as binary vectors).
• Semantic Embeddings: Vectors on a unit sphere.
• Knowledge Graphs: Sets of edges/relations.
• Probabilistic Outputs: Probability simplices (for aleatoric uncertainty).

This mapping induces a push-forward distribution $p(\ell|x)$ over the latent space, derived from the model's predictive distribution $p(s|x)$.

2. Structure-Aware Minimum Bayes Risk (MBR) Decoding

Instead of selecting a single sampled generation, the framework synthesizes a Bayes-optimal response ($\ell_{Bayes}$) directly in the latent space. This is achieved by minimizing the expected risk (Bayes risk) with respect to a task-dependent dissimilarity measure $d_T$:
$$\ell_{Bayes} := \arg \min_{\tilde{\ell} \in \mathcal{L}} \mathbb{E}_{\ell \sim p(\ell|x)} [d_T(\tilde{\ell}, \ell)]$$
Crucially, $\ell_{Bayes}$ does not need to correspond to any specific text generated by the LLM. It is a synthesized action (e.g., a centroid in latent space, a set of majority-voted elements) that minimizes the expected error. For many structures (e.g., sets, graphs, embeddings), this optimal action has a closed-form solution or can be computed efficiently ($O(M)$), unlike raw language MBR which requires $O(M^2)$ pairwise comparisons.

3. Uncertainty Quantification via Bayes Risk

The framework quantifies uncertainty as the Minimum Bayes Risk ($R_{Bayes}$), defined as the risk associated with the optimal action:
$$R_{Bayes} := R(\ell_{Bayes}) = \min_{\tilde{\ell} \in \mathcal{L}} \mathbb{E}_{\ell \sim p(\ell|x)} [d_T(\tilde{\ell}, \ell)]$$
Under the assumption of an unambiguous ground truth (a Dirac delta in latent space), the Bayes risk serves as a lower bound on the true epistemic risk. A high $R_{Bayes}$ provably implies high true risk. This measure captures variability in terms of the task-specific dissimilarity $d_T$, rather than just semantic variability in raw text.

4. Handling Ambiguity

For tasks with inherent ambiguity (aleatoric uncertainty), the framework extends the latent space to the probability simplex. Here, the LLM is prompted to output probability distributions over outcomes. The Bayes-optimal action becomes the mean of these distributions, and the risk is measured via KL divergence, effectively capturing the dispersion of the model's beliefs.

Key Contributions

1. Latent Space Synthesis: The paper introduces a general framework to perform MBR decoding directly in task-specific latent spaces, enabling the synthesis of Bayes-optimal responses that are not constrained to the set of sampled generations.
2. Efficient Closed-Form Solutions: The authors derive closed-form expressions for $\ell_{Bayes}$ and $R_{Bayes}$ across diverse structures (sets, graphs, embeddings, simplices), demonstrating computational efficiency compared to raw language MBR.
3. Task-Aware Uncertainty: They propose using the induced Bayes risk as a principled uncertainty estimator that aligns with task-specific performance metrics, outperforming existing methods that rely on semantic entropy or token-level probabilities.
4. Empirical Validation: The framework is validated across single/multi-answer QA, text summarization (via knowledge graphs), and machine translation, showing consistent improvements in generation quality and uncertainty calibration.

Experimental Results

The authors evaluated the framework on multiple tasks using Gemma 3 and Qwen 3 models (ranging from 4B to 72B parameters).

• Generation Quality: Across all tasks, the Bayes-optimal latent responses ($\ell_{Bayes}$) consistently outperformed standard decoding methods (beam search, greedy, self-consistency) and raw language MBR ($\ell_{Sample}$).
  • In Multi-Answer QA, the method effectively synthesized answers by aggregating information, often outperforming beam search significantly when the latent distribution had high entropy (high uncertainty).
  • In Machine Translation, the method achieved higher COMET scores and lower cosine distances compared to baselines.
  • In Summarization, represented as knowledge graphs, the method improved AlignScore and reduced structural Hamming distance.
• Uncertainty Quantification: The Bayes risk ($R_{Bayes}$) demonstrated superior alignment with actual generation quality compared to baselines like Semantic Entropy (SE), Maximum Sentence Probability (MSP), and Kernel Language Entropy (KLE).
  • Measured via Prediction Rejection Ratio (PRR) and Concordance AUC, the proposed method achieved the best or competitive performance on 4 out of 5 task structures.
  • The results indicate that the method is particularly effective at identifying low-quality generations in complex structural tasks where semantic variability alone is insufficient.

Significance and Claims

The paper claims that task-awareness is a critical, underutilized dimension for improving LLM reliability. By explicitly modeling the latent structure of the response, the framework enables:

1. Superior Decoding: The ability to synthesize optimal responses that are mathematically grounded in the task's error metric, rather than relying on the LLM to generate the perfect text directly.
2. Reliable Uncertainty: A principled uncertainty measure that reflects the actual difficulty of the task (structural dissimilarity) rather than just linguistic variation.

The authors position this as a general paradigm applicable to any problem admitting a latent response structure. They argue that this approach facilitates the reliable integration of LLMs into critical downstream applications where output structure and correctness are paramount. The work suggests that interpreting LLM outputs in a task-aware space extends beyond generation quality, offering a conceptual framework for addressing open problems in LLM research in a principled manner.

The paper acknowledges limitations, noting that the efficiency depends on the tractability of computing the Bayes-optimal action and the feasibility of mapping language to the latent space (and potentially back). However, they assert that for many structured tasks, the concrete latent response is more valuable than the linguistic form, making the framework highly practical.