# Task-Awareness Improves LLM Generations and Uncertainty

arXiv: https://arxiv.org/abs/2601.21500 · HTML: https://mlsanity.com/paper/2601.21500/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 47 references that are in the map.
- First posted on arXiv 29 Jan 2026.

## Builds on (most cited first, 30 of 35)

- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension](https://mlsanity.com/paper/1705.03551/) (2017)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [Get To The Point: Summarization with Pointer-Generator Networks](https://mlsanity.com/paper/1704.04368/) (2017)
- [COMET: A Neural Framework for MT Evaluation](https://mlsanity.com/paper/2009.09025/) (2020)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [Just Ask for Calibration: Strategies for Eliciting Calibrated Confidence Scores from Language Models Fine-Tuned with Human Feedback](https://mlsanity.com/paper/2305.14975/) (2023)
- [Benchmarking Large Language Models for News Summarization](https://mlsanity.com/paper/2301.13848/) (2023)
- [AlignScore: Evaluating Factual Consistency with a Unified Alignment Function](https://mlsanity.com/paper/2305.16739/) (2023)
- [A Contrastive Framework for Neural Text Generation](https://mlsanity.com/paper/2202.06417/) (2022)
- [Kernel Language Entropy: Fine-grained Uncertainty Quantification for LLMs from Semantic Similarities](https://mlsanity.com/paper/2405.20003/) (2024)
- [Shifting Attention to Relevance: Towards the Predictive Uncertainty Quantification of Free-Form Large Language Models](https://mlsanity.com/paper/2307.01379/) (2023)
- [Prompting PaLM for Translation: Assessing Strategies and Performance](https://mlsanity.com/paper/2211.09102/) (2022)
- [Quantifying Aleatoric and Epistemic Uncertainty in Machine Learning: Are Conditional Entropy and Mutual Information Appropriate Measures?](https://mlsanity.com/paper/2209.03302/) (2022)
- [Benchmarking Uncertainty Quantification Methods for Large Language Models with LM-Polygraph](https://mlsanity.com/paper/2406.15627/) (2024)
- [A Comprehensive Survey on Process-Oriented Automatic Text Summarization with Exploration of LLM-Based Methods](https://mlsanity.com/paper/2403.02901/) (2024)
- [Uncertainty Quantification and Confidence Calibration in Large Language Models: A Survey](https://mlsanity.com/paper/2503.15850/) (2025)
- [Sampling-Based Approximations to Minimum Bayes Risk Decoding for Neural Machine Translation](https://mlsanity.com/paper/2108.04718/) (2021)
- [Epsilon Sampling Rocks: Investigating Sampling Strategies for Minimum Bayes Risk Decoding for Machine Translation](https://mlsanity.com/paper/2305.09860/) (2023)
- [Understanding the Properties of Minimum Bayes Risk Decoding in Neural Machine Translation](https://mlsanity.com/paper/2105.08504/) (2021)
- [It's MBR All the Way Down: Modern Generation Techniques Through the Lens of Minimum Bayes Risk](https://mlsanity.com/paper/2310.01387/) (2023)
- [KGGen: Extracting Knowledge Graphs from Plain Text with Language Models](https://mlsanity.com/paper/2502.09956/) (2025)
- [Rethinking Aleatoric and Epistemic Uncertainty](https://mlsanity.com/paper/2412.20892/) (2024)
- [Second-Order Uncertainty Quantification: A Distance-Based Approach](https://mlsanity.com/paper/2312.00995/) (2023)
- [Uncertainty Quantification for LLMs through Minimum Bayes Risk: Bridging Confidence and Consistency](https://mlsanity.com/paper/2502.04964/) (2025)
- [MAQA: Evaluating Uncertainty Quantification in LLMs Regarding Data Uncertainty](https://mlsanity.com/paper/2408.06816/) (2024)
- [LogParser-LLM: Advancing Efficient Log Parsing with Large Language Models](https://mlsanity.com/paper/2408.13727/) (2024)
- [Better Instruction-Following Through Minimum Bayes Risk](https://mlsanity.com/paper/2410.02902/) (2024)
- [The Illusion of Certainty: Uncertainty Quantification for LLMs Fails under Ambiguity](https://mlsanity.com/paper/2511.04418/) (2025)

## Built on by (most cited first, 3 of 3)

- [Decoding as Optimisation on the Probability Simplex: From Top-K to Top-P (Nucleus) to Best-of-K Samplers](https://mlsanity.com/paper/2602.18292/) (2026)
- [Task-Aware Calibration: Provably Optimal Decoding in LLMs](https://mlsanity.com/paper/2605.10202/) (2026)
- [What Is a Clarification? Estimating Model-Dependent Input Ambiguity](https://mlsanity.com/paper/2610.20052/) (2026)

## Abstract (arXiv)

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.