Focus on the Likely: Test-time Instance-based Uncertainty Removal

2025-05-02

AI Summary

This work addresses the challenge of inaccurate predictions under high uncertainty during model inference. We propose an instance-level test-time fine-tuning (TTFT) method that requires no additional training data. Its core is a “focus-on-likely-classes” mechanism: for high-entropy predictions, it performs a single-step, uncertainty-driven gradient update during inference while dynamically suppressing zero-probability classes. To our knowledge, this is the first TTFT paradigm relying solely on the test sample itself—without external supervision or auxiliary data. We theoretically characterize its differential calibration effects on shared versus task-specific features. Empirically, the method significantly improves accuracy on high-uncertainty samples across diverse text and image classification tasks. Remarkably, it achieves consistent performance gains using a single, fixed set of hyperparameters across all tasks, demonstrating both computational efficiency and strong cross-task generalizability.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Inference
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

Abstract

We propose two novel test-time fine-tuning methods to improve uncertain model predictions. Our methods require no auxiliary data and use the given test instance only. Instead of performing a greedy selection of the most likely class to make a prediction, we introduce an additional focus on the likely classes step during inference. By applying a single-step gradient descent, we refine predictions when an initial forward pass indicates high uncertainty. This aligns predictions more closely with the ideal of assigning zero probability to less plausible outcomes. Our theoretical discussion provides a deeper understanding highlighting the impact on shared and non-shared features among (focus) classes. The experimental evaluation highlights accuracy gains on samples exhibiting high decision uncertainty for a diverse set of models from both the text and image domain using the same hyperparameters.

Problem

• Improve uncertain model predictions without auxiliary data
• Refine predictions using single-step gradient descent
• Enhance accuracy for high-uncertainty text and image samples

Innovation

• Test-time fine-tuning without auxiliary data
• Single-step gradient descent for uncertainty refinement
• Focus on likely classes during inference