Title: Focus on Likely Classes for Test-Time Prediction
Primary area: optimization
Keywords: test-time computing, fine-tuning, reflection, reasoning

Abstract: We ask: Does focusing on classes predicted as likely improve model predictions? We aim for an affirmative answer by proposing two novel test-time fine-tuning methods to improve uncertain model predictions. Instead of greedily selecting the most likely class, we introduce an additional step, focus on the likely classes, to refine predictions. By applying a single gradient descent step with a large learning rate, 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. The experimental evaluation demonstrates accuracy gains for one of our methods on average, which emphasizes shared features among likely classes, across diverse text and image domain models.